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

filter on checkpoint NS

Browse Source
This commit is contained in:
vbarda
2024-08-26 19:26:35 -04:00
parent bf4dc5d5f6
commit 85e698e20b
5 changed files with 145 additions and 286 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py
+4
View File
@@ -255,6 +255,10 @@ class BasePostgresSaver(BaseCheckpointSaver):
if config:
wheres.append("thread_id = %s ")
param_values.append(config["configurable"]["thread_id"])
if checkpoint_ns := config["configurable"].get("checkpoint_ns"):
wheres.append("checkpoint_ns = %s")
param_values.append(checkpoint_ns)
if checkpoint_id := get_checkpoint_id(config):
wheres.append("checkpoint_id = %s ")
param_values.append(checkpoint_id)
libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/utils.py
+4
View File
@@ -70,6 +70,10 @@ def search_where(
if config is not None:
wheres.append("thread_id = ?")
param_values.append(config["configurable"]["thread_id"])
if checkpoint_ns := config["configurable"].get("checkpoint_ns"):
wheres.append("checkpoint_ns = ?")
param_values.append(checkpoint_ns)
if checkpoint_id := get_checkpoint_id(config):
wheres.append("checkpoint_id = ?")
param_values.append(checkpoint_id)
libs/checkpoint/langgraph/checkpoint/memory/__init__.py
+4
View File
@@ -177,9 +177,13 @@ class MemorySaver(
Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples.
"""
thread_ids = (config["configurable"]["thread_id"],) if config else self.storage
config_checkpoint_ns = config["configurable"].get("checkpoint_ns") if config else None
config_checkpoint_id = get_checkpoint_id(config) if config else None
for thread_id in thread_ids:
for checkpoint_ns in self.storage[thread_id].keys():
if config_checkpoint_ns and checkpoint_ns != config_checkpoint_ns:
continue
for checkpoint_id, (
checkpoint,
metadata_b,
libs/langgraph/tests/test_pregel.py
+67 -143
View File
@@ -10442,81 +10442,9 @@ def test_nested_graph_state(
}
},
)
# test loading inner snapshot
child_snapshot = app.get_state(
{"configurable": {"thread_id": "1", "checkpoint_ns": "inner"}}
)
assert child_snapshot == StateSnapshot(
values={
"my_key": "hi my value here and there",
"my_other_key": "hi my value here",
},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "inner",
"checkpoint_id": 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",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
# test looking up parent state by checkpoint ID
assert app.get_state(
{
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": child_snapshot.config["configurable"]["checkpoint_id"],
}
},
) == StateSnapshot(
values={"my_key": "hi my value"},
tasks=(PregelTask(AnyStr(), "inner"),),
next=("inner",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={"inner": child_snapshot},
)
# test full history at the end
assert list(app.get_state_history(config)) == [
actual_history = list(app.get_state_history(config))
expected_history = [
StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
tasks=(),
@@ -10675,6 +10603,10 @@ def test_nested_graph_state(
subgraph_state_snapshots=None,
),
]
assert actual_history == expected_history
# test looking up parent state by checkpoint ID
for actual_snapshot, expected_snapshot in zip(actual_history, expected_history):
assert app.get_state(actual_snapshot.config) == expected_snapshot
@pytest.mark.parametrize(
@@ -10847,75 +10779,11 @@ def test_doubly_nested_graph_state(
}
},
)
# test getting snapshot by ID
config = list(app.get_state_history(config))[2].config
# test getting grandchild snapshot
grandchild_snapshot = app.get_state(
{"configurable": {"thread_id": "1", "checkpoint_ns": "child|child_1"}}
)
assert grandchild_snapshot == StateSnapshot(
values={"my_key": "hi my value here and there"},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child|child_1",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"grandchild_2": {"my_key": "hi my value here and there"}},
"step": 2,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child|child_1",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
# test getting child snapshot
child_snapshot = app.get_state(
{"configurable": {"thread_id": "1", "checkpoint_ns": "child"}},
)
assert child_snapshot == StateSnapshot(
values={"my_key": "hi my value here and there"},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"child_1": {"my_key": "hi my value here and there"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={"child_1": grandchild_snapshot},
)
# test getting parent snapshot for a checkpoint ID
assert app.get_state(
{
"configurable": {
"thread_id": "1",
"checkpoint_id": child_snapshot.config["configurable"]["checkpoint_id"],
}
},
) == StateSnapshot(
assert app.get_state(config) == StateSnapshot(
values={"my_key": "hi my value"},
tasks=(PregelTask(AnyStr(), "child"),),
next=("child",),
@@ -10939,7 +10807,63 @@ def test_doubly_nested_graph_state(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={"child": child_snapshot},
subgraph_state_snapshots={
"child": StateSnapshot(
values={"my_key": "hi my value here and there"},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"child_1": {"my_key": "hi my value here and there"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"child_1": StateSnapshot(
values={"my_key": "hi my value here and there"},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child|child_1",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {
"grandchild_2": {"my_key": "hi my value here and there"}
},
"step": 2,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child|child_1",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
)
},
)
libs/langgraph/tests/test_pregel_async.py
+66 -143
View File
@@ -8946,81 +8946,9 @@ async def test_nested_graph_state(
}
},
)
# test loading inner snapshot
child_snapshot = await app.aget_state(
{"configurable": {"thread_id": "1", "checkpoint_ns": "inner"}}
)
assert child_snapshot == StateSnapshot(
values={
"my_key": "hi my value here and there",
"my_other_key": "hi my value here",
},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "inner",
"checkpoint_id": 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",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
# test looking up parent state by checkpoint ID
assert await app.aget_state(
{
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": child_snapshot.config["configurable"]["checkpoint_id"],
}
},
) == StateSnapshot(
values={"my_key": "hi my value"},
tasks=(PregelTask(AnyStr(), "inner"),),
next=("inner",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={"inner": child_snapshot},
)
# test full history at the end
assert [s async for s in app.aget_state_history(config)] == [
actual_history = [s async for s in app.aget_state_history(config)]
expected_history = [
StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
tasks=(),
@@ -9179,6 +9107,10 @@ async def test_nested_graph_state(
subgraph_state_snapshots=None,
),
]
assert actual_history == expected_history
# test looking up parent state by checkpoint ID
for actual_snapshot, expected_snapshot in zip(actual_history, expected_history):
assert await app.aget_state(actual_snapshot.config) == expected_snapshot
@pytest.mark.parametrize(
@@ -9356,75 +9288,10 @@ async def test_doubly_nested_graph_state(
}
},
)
# test getting snapshot by ID
config = [s async for s in app.aget_state_history(config)][2].config
# test getting grandchild snapshot
grandchild_snapshot = await app.aget_state(
{"configurable": {"thread_id": "1", "checkpoint_ns": "child|child_1"}}
)
assert grandchild_snapshot == StateSnapshot(
values={"my_key": "hi my value here and there"},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child|child_1",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"grandchild_2": {"my_key": "hi my value here and there"}},
"step": 2,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child|child_1",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
# test getting child snapshot
child_snapshot = await app.aget_state(
{"configurable": {"thread_id": "1", "checkpoint_ns": "child"}},
)
assert child_snapshot == StateSnapshot(
values={"my_key": "hi my value here and there"},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"child_1": {"my_key": "hi my value here and there"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={"child_1": grandchild_snapshot},
)
# test getting parent snapshot for a checkpoint ID
assert await app.aget_state(
{
"configurable": {
"thread_id": "1",
"checkpoint_id": child_snapshot.config["configurable"]["checkpoint_id"],
}
},
) == StateSnapshot(
assert await app.aget_state(config) == StateSnapshot(
values={"my_key": "hi my value"},
tasks=(PregelTask(AnyStr(), "child"),),
next=("child",),
@@ -9448,7 +9315,63 @@ async def test_doubly_nested_graph_state(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={"child": child_snapshot},
subgraph_state_snapshots={
"child": StateSnapshot(
values={"my_key": "hi my value here and there"},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"child_1": {"my_key": "hi my value here and there"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"child_1": StateSnapshot(
values={"my_key": "hi my value here and there"},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child|child_1",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {
"grandchild_2": {"my_key": "hi my value here and there"}
},
"step": 2,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child|child_1",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
)
},
)