WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-22 07:32:25 +02:00
Code Issues Packages Projects Releases Wiki Activity

cleanup

Browse Source
This commit is contained in:
vbarda
2024-08-12 12:43:28 -04:00
parent c9d6a41d75
commit 322cfc46d3
5 changed files with 1057 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint/langgraph/checkpoint/base/__init__.py
+2
View File
@@ -259,6 +259,7 @@ class BaseCheckpointSaver(ABC):
filter: Optional[Dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
include_nested_checkpoints: bool = False,
) -> Iterator[CheckpointTuple]:
"""List checkpoints that match the given criteria.
@@ -350,6 +351,7 @@ class BaseCheckpointSaver(ABC):
filter: Optional[Dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
include_nested_checkpoints: bool = False,
) -> AsyncIterator[CheckpointTuple]:
"""Asynchronously list checkpoints that match the given criteria.
libs/checkpoint/langgraph/checkpoint/memory/__init__.py
+63 -46
View File
@@ -157,6 +157,7 @@ class MemorySaver(
filter: Optional[Dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
include_nested_checkpoints: bool = False,
) -> Iterator[CheckpointTuple]:
"""List checkpoints from the in-memory storage.
@@ -177,53 +178,67 @@ class MemorySaver(
config["configurable"].get("checkpoint_ns", "") if config else ""
)
for thread_id in thread_ids:
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 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_typed(metadata_b)
if filter and not all(
query_value == metadata[query_key]
for query_key, query_value in filter.items()
):
continue
# limit search results
if limit is not None and limit <= 0:
break
elif limit is not None:
limit -= 1
yield CheckpointTuple(
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,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": parent_checkpoint_id,
}
}
if parent_checkpoint_id
else None,
checkpoint_ns_iter = (
(
key
for key in self.storage[thread_id].keys()
if key.startswith(checkpoint_ns)
)
if include_nested_checkpoints
else [checkpoint_ns]
)
for checkpoint_ns in checkpoint_ns_iter:
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 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_typed(metadata_b)
if filter and not all(
query_value == metadata[query_key]
for query_key, query_value in filter.items()
):
continue
# limit search results
if limit is not None and limit <= 0:
break
elif limit is not None:
limit -= 1
yield CheckpointTuple(
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,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": parent_checkpoint_id,
}
}
if parent_checkpoint_id
else None,
)
def put(
self,
@@ -315,6 +330,7 @@ class MemorySaver(
filter: Optional[Dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
include_nested_checkpoints: bool = False,
) -> AsyncIterator[CheckpointTuple]:
"""Asynchronous version of list.
@@ -335,6 +351,7 @@ class MemorySaver(
before=before,
limit=limit,
filter=filter,
include_nested_checkpoints=include_nested_checkpoints,
),
config,
)
libs/langgraph/langgraph/pregel/__init__.py
+90 -60
View File
@@ -65,12 +65,12 @@ from langgraph.checkpoint.base import (
empty_checkpoint,
)
from langgraph.constants import (
CHECKPOINT_NAMESPACE_SEPARATOR,
CONFIG_KEY_CHECKPOINTER,
CONFIG_KEY_READ,
CONFIG_KEY_RESUMING,
CONFIG_KEY_SEND,
INTERRUPT,
THREAD_ID_SEPARATOR,
)
from langgraph.errors import GraphRecursionError, InvalidUpdateError
from langgraph.managed.base import (
@@ -422,35 +422,36 @@ class Pregel(
@staticmethod
def _assemble_state_snapshot_hierarchy(
root_thread_id: str, thread_id_to_state_snapshots: dict[str, StateSnapshot]
root_checkpoint_ns: str,
checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot],
) -> StateSnapshot:
thread_ids_to_visit = sorted(
thread_id_to_state_snapshots.keys(),
key=lambda x: len(x.split(THREAD_ID_SEPARATOR)),
checkpoint_ns_list_to_visit = sorted(
checkpoint_ns_to_state_snapshots.keys(),
key=lambda x: len(x.split(CHECKPOINT_NAMESPACE_SEPARATOR)),
)
while thread_ids_to_visit:
thread_id = thread_ids_to_visit.pop()
state_snapshot = thread_id_to_state_snapshots[thread_id]
*path, subgraph_node = thread_id.split(THREAD_ID_SEPARATOR)
parent_thread_id = THREAD_ID_SEPARATOR.join(path)
if parent_thread_id and (
parent_state_snapshot := thread_id_to_state_snapshots.get(
parent_thread_id
while checkpoint_ns_list_to_visit:
checkpoint_ns = checkpoint_ns_list_to_visit.pop()
state_snapshot = checkpoint_ns_to_state_snapshots[checkpoint_ns]
*path, subgraph_node = checkpoint_ns.split(CHECKPOINT_NAMESPACE_SEPARATOR)
parent_checkpoint_ns = CHECKPOINT_NAMESPACE_SEPARATOR.join(path)
if subgraph_node and (
parent_state_snapshot := checkpoint_ns_to_state_snapshots.get(
parent_checkpoint_ns
)
):
parent_subgraph_snapshots = {
**(parent_state_snapshot.subgraph_state_snapshots or {}),
subgraph_node: state_snapshot,
}
thread_id_to_state_snapshots[
parent_thread_id
] = thread_id_to_state_snapshots[parent_thread_id]._replace(
checkpoint_ns_to_state_snapshots[
parent_checkpoint_ns
] = checkpoint_ns_to_state_snapshots[parent_checkpoint_ns]._replace(
subgraph_state_snapshots=parent_subgraph_snapshots
)
state_snapshot = thread_id_to_state_snapshots.pop(root_thread_id, None)
state_snapshot = checkpoint_ns_to_state_snapshots.pop(root_checkpoint_ns, None)
if state_snapshot is None:
raise ValueError(f"Missing checkpoint for thread ID '{root_thread_id}'")
raise ValueError(f"Missing checkpoint for thread ID '{root_checkpoint_ns}'")
return state_snapshot
def get_state(
@@ -461,36 +462,51 @@ class Pregel(
raise ValueError("No checkpointer set")
if include_subgraph_state:
checkpoint_tuples = self.checkpointer.list(config, as_prefix=True)
checkpoint_tuples = self.checkpointer.list(
config, include_nested_checkpoints=True
)
else:
checkpoint_tuples = iter([self.checkpointer.get_tuple(config)])
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] = {}
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = config["configurable"].get("checkpoint_id")
checkpoint_ns_to_checkpoint_id: dict[str, str] = {}
checkpoint_ns_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:
saved_checkpoint_ns = checkpoint_tuple.config["configurable"][
"checkpoint_ns"
]
saved_checkpoint_id = checkpoint_tuple.config["configurable"][
"checkpoint_id"
]
if checkpoint_id and checkpoint_id != saved_checkpoint_id:
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:
existing_checkpoint_id = checkpoint_ns_to_checkpoint_id.get(
saved_checkpoint_ns
)
# keep only most recent checkpoint_id
if (
existing_checkpoint_id is None
or saved_checkpoint_id > existing_checkpoint_id
):
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
checkpoint_ns_to_state_snapshots[saved_checkpoint_ns] = state_snapshot
checkpoint_ns_to_checkpoint_id[
saved_checkpoint_ns
] = saved_checkpoint_id
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}'"
if not checkpoint_ns_to_state_snapshots:
error_msg = (
f"Could not find checkpoints for checkpoint NS '{checkpoint_ns}'"
)
if checkpoint_id:
error_msg += f" and checkpoint ID '{checkpoint_id}'"
raise ValueError(error_msg)
state_snapshot = self._assemble_state_snapshot_hierarchy(
thread_id, thread_id_to_state_snapshots
checkpoint_ns, checkpoint_ns_to_state_snapshots
)
return state_snapshot
@@ -502,7 +518,9 @@ class Pregel(
raise ValueError("No checkpointer set")
if include_subgraph_state:
checkpoint_tuples = self.checkpointer.alist(config, as_prefix=True)
checkpoint_tuples = self.checkpointer.alist(
config, include_nested_checkpoints=True
)
else:
async def alist_checkpoints():
@@ -510,34 +528,45 @@ class Pregel(
checkpoint_tuples = alist_checkpoints()
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] = {}
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = config["configurable"].get("checkpoint_id")
checkpoint_ns_to_checkpoint_id: dict[str, str] = {}
checkpoint_ns_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:
saved_checkpoint_ns = checkpoint_tuple.config["configurable"][
"checkpoint_ns"
]
saved_checkpoint_id = checkpoint_tuple.config["configurable"][
"checkpoint_id"
]
if checkpoint_id and checkpoint_id != saved_checkpoint_id:
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
existing_checkpoint_id = checkpoint_ns_to_checkpoint_id.get(
saved_checkpoint_ns
)
# keep only most recent checkpoint_id
if (
existing_checkpoint_id is None
or saved_checkpoint_id > existing_checkpoint_id
):
state_snapshot = self._prepare_state_snapshot(checkpoint_tuple, config)
checkpoint_ns_to_state_snapshots[saved_checkpoint_ns] = state_snapshot
checkpoint_ns_to_checkpoint_id[
saved_checkpoint_ns
] = saved_checkpoint_id
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}'"
if not checkpoint_ns_to_state_snapshots:
error_msg = (
f"Could not find checkpoints for checkpoint NS '{checkpoint_ns}'"
)
if checkpoint_id:
error_msg += f" and checkpoint ID '{checkpoint_id}'"
raise ValueError(error_msg)
state_snapshot = self._assemble_state_snapshot_hierarchy(
thread_id, thread_id_to_state_snapshots
checkpoint_ns, checkpoint_ns_to_state_snapshots
)
return state_snapshot
@@ -584,7 +613,7 @@ class Pregel(
-1,
for_execution=False,
)
yield StateSnapshot(
read_channels(channels, self.stream_channels_asis),
tuple(t.name for t in next_tasks),
@@ -620,10 +649,11 @@ class Pregel(
) in self.checkpointer.alist(config, before=before, limit=limit, filter=filter):
if include_subgraph_state:
state_snapshot = await self.aget_state(
config, include_subgraph_state=True)
config, include_subgraph_state=True
)
yield state_snapshot
else:
async with AsyncChannelsManager(
async with AsyncChannelsManager(
{
k: LastValue(None) if isinstance(c, Context) else c
for k, c in self.channels.items()
libs/langgraph/tests/test_pregel.py
+450
View File
@@ -9093,6 +9093,456 @@ def test_doubly_nested_graph_interrupts(
]
@pytest.mark.parametrize(
"checkpointer_name",
["memory", "sqlite", "postgres", "postgres_pipe"],
)
def test_nested_graph_state(
request: pytest.FixtureRequest, checkpointer_name: str
) -> None:
checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name)
class InnerState(TypedDict):
my_key: str
my_other_key: str
def inner_1(state: InnerState):
return {
"my_key": state["my_key"] + " here",
"my_other_key": state["my_key"],
}
def inner_2(state: InnerState):
return {
"my_key": state["my_key"] + " and there",
"my_other_key": state["my_key"],
}
inner = StateGraph(InnerState)
inner.add_node("inner_1", inner_1)
inner.add_node("inner_2", inner_2)
inner.add_edge("inner_1", "inner_2")
inner.set_entry_point("inner_1")
inner.set_finish_point("inner_2")
class State(TypedDict):
my_key: str
def outer_1(state: State):
return {"my_key": "hi " + state["my_key"]}
def outer_2(state: State):
return {"my_key": state["my_key"] + " and back again"}
graph = StateGraph(State)
graph.add_node("outer_1", outer_1)
graph.add_node("inner", inner.compile(interrupt_before=["inner_2"]))
graph.add_node("outer_2", outer_2)
graph.set_entry_point("outer_1")
graph.add_edge("outer_1", "inner")
graph.add_edge("inner", "outer_2")
graph.set_finish_point("outer_2")
app = graph.compile(checkpointer=checkpointer)
config = {"configurable": {"thread_id": "1"}}
app.invoke({"my_key": "my value"}, config, debug=True)
# test state w/ nested subgraph state (right after interrupt)
assert app.get_state(config, include_subgraph_state=False) == StateSnapshot(
values={"my_key": "hi my value"},
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=None,
)
assert app.get_state(config, include_subgraph_state=True) == StateSnapshot(
values={"my_key": "hi my value"},
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": StateSnapshot(
values={"my_key": "hi my value here"},
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "inner",
"checkpoint_id": 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",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
)
assert list(app.get_state_history(config, include_subgraph_state=True)) == [
StateSnapshot(
values={"my_key": "hi my value"},
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": StateSnapshot(
values={"my_key": "hi my value here"},
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "inner",
"checkpoint_id": 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",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
),
StateSnapshot(
values={},
next=("__start__",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
subgraph_state_snapshots=None,
),
]
app.invoke(None, config, debug=True)
# test state w/ nested subgraph state (after resuming from interrupt)
assert app.get_state(config, include_subgraph_state=True) == StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": 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",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={"my_key": "hi my value here and there"},
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,
)
},
)
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=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": 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",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
),
StateSnapshot(
values={"my_key": "hi my value here and there"},
next=("outer_2",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": 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",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
),
StateSnapshot(
values={"my_key": "hi my value"},
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(),
}
},
# 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 checkpoint_id
# 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",
"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,
)
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
),
StateSnapshot(
values={},
next=("__start__",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
subgraph_state_snapshots=None,
),
]
def test_repeat_condition(snapshot: SnapshotAssertion) -> None:
class AgentState(TypedDict):
hello: str
libs/langgraph/tests/test_pregel_async.py
+452
View File
@@ -7600,6 +7600,458 @@ async def test_doubly_nested_graph_interrupts(
]
@pytest.mark.parametrize(
"checkpointer_name",
["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"],
)
async def test_nested_graph_state(
request: pytest.FixtureRequest, checkpointer_name: str
) -> None:
checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name)
class InnerState(TypedDict):
my_key: str
my_other_key: str
def inner_1(state: InnerState):
return {
"my_key": state["my_key"] + " here",
"my_other_key": state["my_key"],
}
def inner_2(state: InnerState):
return {
"my_key": state["my_key"] + " and there",
"my_other_key": state["my_key"],
}
inner = StateGraph(InnerState)
inner.add_node("inner_1", inner_1)
inner.add_node("inner_2", inner_2)
inner.add_edge("inner_1", "inner_2")
inner.set_entry_point("inner_1")
inner.set_finish_point("inner_2")
class State(TypedDict):
my_key: str
def outer_1(state: State):
return {"my_key": "hi " + state["my_key"]}
def outer_2(state: State):
return {"my_key": state["my_key"] + " and back again"}
graph = StateGraph(State)
graph.add_node("outer_1", outer_1)
graph.add_node("inner", inner.compile(interrupt_before=["inner_2"]))
graph.add_node("outer_2", outer_2)
graph.set_entry_point("outer_1")
graph.add_edge("outer_1", "inner")
graph.add_edge("inner", "outer_2")
graph.set_finish_point("outer_2")
app = graph.compile(checkpointer=checkpointer)
config = {"configurable": {"thread_id": "1"}}
await app.ainvoke({"my_key": "my value"}, config, debug=True)
# test state w/ nested subgraph state (right after interrupt)
assert await app.aget_state(config, include_subgraph_state=False) == StateSnapshot(
values={"my_key": "hi my value"},
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=None,
)
assert app.get_state(config, include_subgraph_state=True) == StateSnapshot(
values={"my_key": "hi my value"},
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": StateSnapshot(
values={"my_key": "hi my value here"},
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "inner",
"checkpoint_id": 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",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
)
assert list(app.get_state_history(config, include_subgraph_state=True)) == [
StateSnapshot(
values={"my_key": "hi my value"},
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": StateSnapshot(
values={"my_key": "hi my value here"},
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "inner",
"checkpoint_id": 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",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
),
StateSnapshot(
values={},
next=("__start__",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": 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(config, include_subgraph_state=True) == StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": 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",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={"my_key": "hi my value here and there"},
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,
)
},
)
assert [
s async for s 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",
"checkpoint_ns": "",
"checkpoint_id": 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",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
),
StateSnapshot(
values={"my_key": "hi my value here and there"},
next=("outer_2",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": 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",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
),
StateSnapshot(
values={"my_key": "hi my value"},
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(),
}
},
# 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 checkpoint_id
# 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",
"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,
)
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
),
StateSnapshot(
values={},
next=("__start__",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
subgraph_state_snapshots=None,
),
]
async def test_checkpoint_metadata() -> None:
"""This test verifies that a run's configurable fields are merged with the
previous checkpoint config for each step in the run.