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

use .list for looking up prefix-matched checkpoints

Browse Source
This commit is contained in:
vbarda
2024-07-24 13:45:48 -04:00
parent 6ae2c6c7b9
commit a7d48465da
7 changed files with 333 additions and 205 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/checkpoint/aiosqlite.py
+2 -45
View File
@@ -306,6 +306,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
filter: Optional[Dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
as_prefix: bool = False,
) -> AsyncIterator[CheckpointTuple]:
"""List checkpoints from the database asynchronously.
@@ -322,7 +323,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
AsyncIterator[CheckpointTuple]: An asynchronous iterator of matching checkpoint tuples.
"""
await self.setup()
where, param_values = search_where(config, filter, before)
where, param_values = search_where(config, filter, before, as_prefix=as_prefix)
query = f"""SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata
FROM checkpoints
{where}
@@ -347,50 +348,6 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager):
),
)
async def alist_subgraph_checkpoints(
self, config: RunnableConfig
) -> AsyncIterator[CheckpointTuple]:
# TODO: docstring
async with self.conn.cursor() as cur:
if config["configurable"].get("thread_ts"):
await cur.execute(
"SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id LIKE ? || '%' AND thread_ts = ?",
(
str(config["configurable"]["thread_id"]),
str(config["configurable"]["thread_ts"]),
),
)
else:
await cur.execute(
"""SELECT checkpoints.thread_id, checkpoints.thread_ts, checkpoints.parent_ts, checkpoints.checkpoint, checkpoints.metadata
FROM checkpoints
INNER JOIN (
SELECT thread_id, MAX(thread_ts) as thread_ts
FROM checkpoints
WHERE thread_id LIKE ? || '%'
GROUP BY thread_id
) latest_checkpoints
ON checkpoints.thread_id = latest_checkpoints.thread_id AND checkpoints.thread_ts = latest_checkpoints.thread_ts
ORDER BY checkpoints.thread_id, checkpoints.thread_ts DESC""",
(str(config["configurable"]["thread_id"]),),
)
async for thread_id, thread_ts, parent_ts, value, metadata in cur:
yield CheckpointTuple(
{"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}},
self.serde.loads(value),
self.serde.loads(metadata) if metadata is not None else {},
(
{
"configurable": {
"thread_id": thread_id,
"thread_ts": parent_ts,
}
}
if parent_ts
else None
),
)
async def aput(
self,
config: RunnableConfig,
libs/langgraph/langgraph/checkpoint/base.py
+1 -12
View File
@@ -214,6 +214,7 @@ class BaseCheckpointSaver(ABC):
filter: Optional[Dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
as_prefix: bool = False,
) -> Iterator[CheckpointTuple]:
"""List checkpoints that match the given criteria.
@@ -231,12 +232,6 @@ class BaseCheckpointSaver(ABC):
"""
raise NotImplementedError
def list_subgraph_checkpoints(
self, config: RunnableConfig
) -> Iterator[CheckpointTuple]:
# TODO: docstring
raise NotImplementedError
def put(
self,
config: RunnableConfig,
@@ -329,12 +324,6 @@ class BaseCheckpointSaver(ABC):
raise NotImplementedError
yield
async def alist_subgraph_checkpoints(
self, config: RunnableConfig
) -> AsyncIterator[CheckpointTuple]:
# TODO: docstring
raise NotImplementedError
async def aput(
self,
config: RunnableConfig,
libs/langgraph/langgraph/checkpoint/memory.py
+24 -50
View File
@@ -119,6 +119,7 @@ class MemorySaver(BaseCheckpointSaver):
filter: Optional[Dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
as_prefix: bool = False,
) -> Iterator[CheckpointTuple]:
"""List checkpoints from the in-memory storage.
@@ -134,7 +135,19 @@ class MemorySaver(BaseCheckpointSaver):
Yields:
Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples.
"""
thread_ids = (config["configurable"]["thread_id"],) if config else self.storage
if config:
config_thread_id = config["configurable"]["thread_id"]
if as_prefix:
thread_ids = (
thread_id
for thread_id in self.storage
if thread_id.startswith(config_thread_id)
)
else:
thread_ids = (config_thread_id,)
else:
thread_ids = self.storage.keys()
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
@@ -171,40 +184,6 @@ class MemorySaver(BaseCheckpointSaver):
else None,
)
def list_subgraph_checkpoints(
self, config: RunnableConfig
) -> Iterator[CheckpointTuple]:
# TODO: docstring
thread_id_prefix = config["configurable"]["thread_id"]
matching_thread_ids = [
key for key in self.storage.keys() if key.startswith(thread_id_prefix)
]
for thread_id in matching_thread_ids:
ts = config["configurable"].get("thread_ts")
if not ts:
if checkpoints := self.storage[thread_id]:
ts = max(checkpoints.keys())
if saved := self.storage[thread_id].get(ts):
checkpoint, metadata, parent_ts = saved
writes = self.writes[(thread_id, ts)]
yield CheckpointTuple(
config={"configurable": {"thread_id": thread_id, "thread_ts": ts}},
checkpoint=self.serde.loads(checkpoint),
metadata=self.serde.loads(metadata),
pending_writes=[
(id, c, self.serde.loads(v)) for id, c, v in writes
],
parent_config={
"configurable": {
"thread_id": thread_id,
"thread_ts": parent_ts,
}
}
if parent_ts
else None,
)
def put(
self,
config: RunnableConfig,
@@ -288,6 +267,7 @@ class MemorySaver(BaseCheckpointSaver):
filter: Optional[Dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
as_prefix: bool = False,
) -> AsyncIterator[CheckpointTuple]:
"""Asynchronous version of list.
@@ -302,7 +282,15 @@ class MemorySaver(BaseCheckpointSaver):
"""
loop = asyncio.get_running_loop()
iter = await loop.run_in_executor(
None, partial(self.list, before=before, limit=limit, filter=filter), config
None,
partial(
self.list,
before=before,
limit=limit,
filter=filter,
as_prefix=as_prefix,
),
config,
)
while True:
# handling StopIteration exception inside coroutine won't work
@@ -312,20 +300,6 @@ class MemorySaver(BaseCheckpointSaver):
else:
break
async def alist_subgraph_checkpoints(
self, config: RunnableConfig
) -> AsyncIterator[CheckpointTuple]:
# TODO: docstring
loop = asyncio.get_running_loop()
iter = await loop.run_in_executor(None, self.list_subgraph_checkpoints, config)
while True:
# handling StopIteration exception inside coroutine won't work
# as expected, so using next() with default value to break the loop
if item := await loop.run_in_executor(None, next, iter, None):
yield item
else:
break
async def aput(
self,
config: RunnableConfig,
libs/langgraph/langgraph/checkpoint/sqlite.py
+6 -46
View File
@@ -301,6 +301,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
filter: Optional[Dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
as_prefix: bool = False,
) -> Iterator[CheckpointTuple]:
"""List checkpoints from the database.
@@ -331,7 +332,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
>>> print(checkpoints)
[CheckpointTuple(...), ...]
"""
where, param_values = search_where(config, filter, before)
where, param_values = search_where(config, filter, before, as_prefix=as_prefix)
query = f"""SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata
FROM checkpoints
{where}
@@ -357,50 +358,6 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager):
),
)
def list_subgraph_checkpoints(
self, config: RunnableConfig
) -> Iterator[CheckpointTuple]:
# TODO: docstring
with self.cursor(transaction=False) as cur:
if config["configurable"].get("thread_ts"):
cur.execute(
"SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id LIKE ? || '%' AND thread_ts = ?",
(
str(config["configurable"]["thread_id"]),
str(config["configurable"]["thread_ts"]),
),
)
else:
cur.execute(
"""SELECT checkpoints.thread_id, checkpoints.thread_ts, checkpoints.parent_ts, checkpoints.checkpoint, checkpoints.metadata
FROM checkpoints
INNER JOIN (
SELECT thread_id, MAX(thread_ts) as thread_ts
FROM checkpoints
WHERE thread_id LIKE ? || '%'
GROUP BY thread_id
) latest_checkpoints
ON checkpoints.thread_id = latest_checkpoints.thread_id AND checkpoints.thread_ts = latest_checkpoints.thread_ts
ORDER BY checkpoints.thread_id, checkpoints.thread_ts DESC""",
(str(config["configurable"]["thread_id"]),),
)
for thread_id, thread_ts, parent_ts, value, metadata in cur:
yield CheckpointTuple(
{"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}},
self.serde.loads(value),
self.serde.loads(metadata) if metadata is not None else {},
(
{
"configurable": {
"thread_id": thread_id,
"thread_ts": parent_ts,
}
}
if parent_ts
else None
),
)
def put(
self,
config: RunnableConfig,
@@ -592,6 +549,7 @@ def search_where(
config: Optional[RunnableConfig],
filter: Optional[Dict[str, Any]],
before: Optional[RunnableConfig] = None,
as_prefix: bool = False,
) -> Tuple[str, Sequence[Any]]:
"""Return WHERE clause predicates for (a)search() given metadata filter
and `before` config.
@@ -606,7 +564,9 @@ def search_where(
# construct predicate for config filter
if config is not None:
wheres.append("thread_id = ?")
thread_filter = "thread_id LIKE ? || '%'" if as_prefix else "thread_id = ?"
wheres.append(thread_filter)
param_values.append(config["configurable"]["thread_id"])
# construct predicate for metadata filter
libs/langgraph/langgraph/pregel/__init__.py
+80 -51
View File
@@ -438,7 +438,7 @@ class Pregel(
state_snapshot = thread_id_to_state_snapshots.pop(root_thread_id, None)
if state_snapshot is None:
raise ValueError(f"Missing snapshot for thread ID '{root_thread_id}'")
raise ValueError(f"Missing checkpoint for thread ID '{root_thread_id}'")
return state_snapshot
def get_state(
@@ -449,18 +449,34 @@ class Pregel(
raise ValueError("No checkpointer set")
if include_subgraph_state:
checkpoint_tuples = self.checkpointer.list_subgraph_checkpoints(config)
checkpoint_tuples = self.checkpointer.list(config, as_prefix=True)
else:
checkpoint_tuples = iter([self.checkpointer.get_tuple(config)])
thread_id_to_state_snapshots: dict[str, StateSnapshot] = {
checkpoint.config["configurable"][
"thread_id"
]: self._prepare_state_snapshot(checkpoint, config)
for checkpoint in checkpoint_tuples
}
thread_id = config["configurable"]["thread_id"]
thread_ts = config["configurable"].get("thread_ts")
thread_id_to_thread_ts: dict[str, str] = {}
thread_id_to_state_snapshots: dict[str, StateSnapshot] = {}
for checkpoint_tuple in checkpoint_tuples:
checkpoint_thread_id = checkpoint_tuple.config["configurable"]["thread_id"]
checkpoint_thread_ts = checkpoint_tuple.config["configurable"]["thread_ts"]
if thread_ts and thread_ts != checkpoint_thread_ts:
continue
existing_thread_ts = thread_id_to_thread_ts.get(checkpoint_thread_id)
# keep only most recent thread_ts
if existing_thread_ts is None or checkpoint_thread_ts > existing_thread_ts:
state_snapshot = self._prepare_state_snapshot(checkpoint_tuple, config)
thread_id_to_state_snapshots[checkpoint_thread_id] = state_snapshot
thread_id_to_thread_ts[checkpoint_thread_id] = checkpoint_thread_ts
if not thread_id_to_state_snapshots:
error_msg = f"Could not find checkpoints for thread ID '{thread_id}'"
if thread_ts:
error_msg += f" and thread TS '{thread_ts}'"
raise ValueError(error_msg)
state_snapshot = self._assemble_state_snapshot_hierarchy(
thread_id, thread_id_to_state_snapshots
)
@@ -474,7 +490,7 @@ class Pregel(
raise ValueError("No checkpointer set")
if include_subgraph_state:
checkpoint_tuples = self.checkpointer.alist_subgraph_checkpoints(config)
checkpoint_tuples = self.checkpointer.alist(config, as_prefix=True)
else:
async def alist_checkpoints():
@@ -482,14 +498,32 @@ class Pregel(
checkpoint_tuples = alist_checkpoints()
thread_id_to_state_snapshots: dict[str, StateSnapshot] = {
checkpoint.config["configurable"][
"thread_id"
]: await self._prepare_state_snapshot_async(checkpoint, config)
async for checkpoint in checkpoint_tuples
}
thread_id = config["configurable"]["thread_id"]
thread_ts = config["configurable"].get("thread_ts")
thread_id_to_thread_ts: dict[str, str] = {}
thread_id_to_state_snapshots: dict[str, StateSnapshot] = {}
async for checkpoint_tuple in checkpoint_tuples:
checkpoint_thread_id = checkpoint_tuple.config["configurable"]["thread_id"]
checkpoint_thread_ts = checkpoint_tuple.config["configurable"]["thread_ts"]
if thread_ts and thread_ts != checkpoint_thread_ts:
continue
existing_thread_ts = thread_id_to_thread_ts.get(checkpoint_thread_id)
# keep only most recent thread_ts
if existing_thread_ts is None or checkpoint_thread_ts > existing_thread_ts:
state_snapshot = await self._prepare_state_snapshot_async(
checkpoint_tuple, config
)
thread_id_to_state_snapshots[checkpoint_thread_id] = state_snapshot
thread_id_to_thread_ts[checkpoint_thread_id] = checkpoint_thread_ts
if not thread_id_to_state_snapshots:
error_msg = f"Could not find checkpoints for thread ID '{thread_id}'"
if thread_ts:
error_msg += f" and thread TS '{thread_ts}'"
raise ValueError(error_msg)
state_snapshot = self._assemble_state_snapshot_hierarchy(
thread_id, thread_id_to_state_snapshots
)
@@ -515,20 +549,8 @@ class Pregel(
for config, checkpoint, metadata, parent_config, _ in self.checkpointer.list(
config, before=before, limit=limit, filter=filter
):
# is there a way to do this more efficiently?
if include_subgraph_state:
checkpoint_tuples = self.checkpointer.list_subgraph_checkpoints(config)
thread_id_to_state_snapshots: dict[str, StateSnapshot] = {
checkpoint.config["configurable"][
"thread_id"
]: self._prepare_state_snapshot(checkpoint, config)
for checkpoint in checkpoint_tuples
}
thread_id = config["configurable"]["thread_id"]
state_snapshot = self._assemble_state_snapshot_hierarchy(
thread_id, thread_id_to_state_snapshots
)
state_snapshot = self.get_state(config, include_subgraph_state=True)
yield state_snapshot
else:
with ChannelsManager(
@@ -561,6 +583,7 @@ class Pregel(
filter: Optional[Dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
include_subgraph_state: bool = False,
) -> AsyncIterator[StateSnapshot]:
"""Get the history of the state of the graph."""
if not self.checkpointer:
@@ -577,28 +600,34 @@ class Pregel(
parent_config,
_,
) in self.checkpointer.alist(config, before=before, limit=limit, filter=filter):
async with AsyncChannelsManager(
self.channels, checkpoint, config
) as channels, AsyncManagedValuesManager(
self.managed_values_dict, ensure_config(config), self
) as managed:
next_tasks = prepare_next_tasks(
checkpoint,
self.nodes,
channels,
managed,
config,
-1,
for_execution=False,
)
yield StateSnapshot(
read_channels(channels, self.stream_channels_asis),
tuple(name for name, _ in next_tasks),
config,
metadata,
checkpoint["ts"],
parent_config,
if include_subgraph_state:
state_snapshot = await self.aget_state(
config, include_subgraph_state=True
)
yield state_snapshot
else:
async with AsyncChannelsManager(
self.channels, checkpoint, config
) as channels, AsyncManagedValuesManager(
self.managed_values_dict, ensure_config(config), self
) as managed:
next_tasks = prepare_next_tasks(
checkpoint,
self.nodes,
channels,
managed,
config,
-1,
for_execution=False,
)
yield StateSnapshot(
read_channels(channels, self.stream_channels_asis),
tuple(name for name, _ in next_tasks),
config,
metadata,
checkpoint["ts"],
parent_config,
)
def update_state(
self,
libs/langgraph/tests/test_pregel.py
+1 -1
View File
@@ -8896,7 +8896,7 @@ def test_nested_graph_state(
)
},
)
assert list(app.get_state_history(config, include_subgraph_state=False)) == [
assert list(app.get_state_history(config, include_subgraph_state=True)) == [
StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
next=(),
libs/langgraph/tests/test_pregel_async.py
+219
View File
@@ -7322,6 +7322,78 @@ async def test_nested_graph_state(checkpointer: BaseCheckpointSaver) -> None:
)
},
)
assert [
c async for c in app.aget_state_history(config, include_subgraph_state=True)
] == [
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "1", "thread_ts": AnyStr()}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={"my_key": "hi my value here"},
next=(),
config={
"configurable": {
"thread_id": "1__inner",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {
"inner_1": {
"my_key": "hi my value here",
"my_other_key": "hi my value",
}
},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1__inner",
"thread_ts": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "1", "thread_ts": AnyStr()}
},
subgraph_state_snapshots=None,
),
StateSnapshot(
values={},
next=("__start__",),
config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
subgraph_state_snapshots=None,
),
]
await app.ainvoke(None, config, debug=True)
# test state w/ nested subgraph state (after resuming from interrupt)
assert await app.aget_state(
@@ -7364,6 +7436,153 @@ async def test_nested_graph_state(checkpointer: BaseCheckpointSaver) -> None:
)
},
)
assert [
c async for c in app.aget_state_history(config, include_subgraph_state=True)
] == [
StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
next=(),
config={
"configurable": {
"thread_id": "1",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {
"outer_2": {
"my_key": "hi my value here and there and back again"
}
},
"step": 3,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"thread_ts": AnyStr(),
}
},
subgraph_state_snapshots=None,
),
StateSnapshot(
values={"my_key": "hi my value here and there"},
next=("outer_2",),
config={
"configurable": {
"thread_id": "1",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"inner": {"my_key": "hi my value here and there"}},
"step": 2,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"thread_ts": AnyStr(),
}
},
subgraph_state_snapshots=None,
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={
"configurable": {
"thread_id": "1",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"thread_ts": AnyStr(),
}
},
# TODO: this is likely very confusing for an end user, and we'll probably need to update this.
# right now this is happening due to us overwriting the
# subgraph snapshot after we finish the graph with while the thread_ts
# is the same as when we interrupted
subgraph_state_snapshots={
"inner": StateSnapshot(
values={"my_key": "hi my value here and there"},
next=(),
config={
"configurable": {
"thread_id": "1__inner",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {
"inner_2": {
"my_key": "hi my value here and there",
"my_other_key": "hi my value here",
}
},
"step": 2,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1__inner",
"thread_ts": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={
"configurable": {
"thread_id": "1",
"thread_ts": AnyStr(),
}
},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"thread_ts": AnyStr(),
}
},
subgraph_state_snapshots=None,
),
StateSnapshot(
values={},
next=("__start__",),
config={
"configurable": {
"thread_id": "1",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
subgraph_state_snapshots=None,
),
]
finally:
if hasattr(checkpointer, "__aexit__"):
await checkpointer.__aexit__(None, None, None)