remove include_subgraph_state kwarg

This commit is contained in:
vbarda
2024-08-14 21:07:56 -04:00
parent 7fa97898aa
commit 45054df71a
3 changed files with 862 additions and 182 deletions
+8 -53
View File
@@ -506,26 +506,19 @@ class Pregel(
yield runnable
yield from runnable.subgraphs
def get_state(
self, config: RunnableConfig, *, include_subgraph_state: bool = False
) -> StateSnapshot:
def get_state(self, config: RunnableConfig) -> StateSnapshot:
"""Get the current state of the graph."""
if not self.checkpointer:
raise ValueError("No checkpointer set")
checkpoint_tuple = self.checkpointer.get_tuple(config)
if include_subgraph_state:
checkpoint_tuples = self.checkpointer.list(config)
else:
checkpoint_tuples = iter([checkpoint_tuple] if checkpoint_tuple else [])
checkpoint_config = checkpoint_tuple.config if checkpoint_tuple else config
checkpoint_ns = checkpoint_config["configurable"].get("checkpoint_ns", "")
checkpoint_id = checkpoint_config["configurable"].get("checkpoint_id")
checkpoint_ns_to_checkpoint_id: dict[str, str] = {}
checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot] = {}
checkpoint_ns_to_graph: dict[str, Pregel] = {}
for checkpoint_tuple in checkpoint_tuples:
for checkpoint_tuple in self.checkpointer.list(config):
saved_checkpoint_ns = checkpoint_tuple.config["configurable"][
"checkpoint_ns"
]
@@ -567,31 +560,19 @@ class Pregel(
)
return state_snapshot
async def aget_state(
self, config: RunnableConfig, *, include_subgraph_state: bool = False
) -> StateSnapshot:
async def aget_state(self, config: RunnableConfig) -> StateSnapshot:
"""Get the current state of the graph."""
if not self.checkpointer:
raise ValueError("No checkpointer set")
checkpoint_tuple = await self.checkpointer.aget_tuple(config)
if include_subgraph_state:
checkpoint_tuples = self.checkpointer.alist(config)
else:
async def alist_checkpoints():
if checkpoint_tuple:
yield checkpoint_tuple
checkpoint_tuples = alist_checkpoints()
checkpoint_config = checkpoint_tuple.config if checkpoint_tuple else config
checkpoint_ns = checkpoint_config["configurable"].get("checkpoint_ns", "")
checkpoint_id = checkpoint_config["configurable"].get("checkpoint_id")
checkpoint_ns_to_checkpoint_id: dict[str, str] = {}
checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot] = {}
checkpoint_ns_to_graph: dict[str, Pregel] = {}
async for checkpoint_tuple in checkpoint_tuples:
async for checkpoint_tuple in self.checkpointer.alist(config):
saved_checkpoint_ns = checkpoint_tuple.config["configurable"][
"checkpoint_ns"
]
@@ -640,7 +621,6 @@ class Pregel(
filter: Optional[Dict[str, Any]] = None,
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
include_subgraph_state: bool = False,
) -> Iterator[StateSnapshot]:
"""Get the history of the state of the graph."""
if not self.checkpointer:
@@ -662,20 +642,8 @@ class Pregel(
# only list root checkpoints here
continue
if include_subgraph_state:
state_snapshot = self.get_state(
checkpoint_tuple.config, include_subgraph_state=True
)
yield state_snapshot
else:
graph = _get_subgraph(
self,
checkpoint_tuple.config["configurable"]["checkpoint_ns"],
)
yield _prepare_state_snapshot(
checkpoint_tuple,
graph,
)
state_snapshot = self.get_state(checkpoint_tuple.config)
yield state_snapshot
async def aget_state_history(
self,
@@ -684,7 +652,6 @@ 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:
@@ -706,20 +673,8 @@ class Pregel(
# only list root checkpoints here
continue
if include_subgraph_state:
state_snapshot = await self.aget_state(
checkpoint_tuple.config, include_subgraph_state=True
)
yield state_snapshot
else:
graph = _get_subgraph(
self,
checkpoint_tuple.config["configurable"]["checkpoint_ns"],
)
yield await _prepare_state_snapshot_async(
checkpoint_tuple,
graph,
)
state_snapshot = await self.aget_state(checkpoint_tuple.config)
yield state_snapshot
def update_state(
self,
+427 -62
View File
@@ -7783,6 +7783,41 @@ def test_nested_graph_interrupts(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={
"my_key": "hi my value here",
"my_other_key": "hi my value",
},
next=("inner_2",),
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"},
@@ -7900,6 +7935,41 @@ def test_nested_graph_interrupts(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={
"my_key": "hi my value here and there",
"my_other_key": "hi my value here",
},
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"},
@@ -8071,6 +8141,41 @@ def test_nested_graph_interrupts(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={
"my_key": "hi my value here",
"my_other_key": "hi my value",
},
next=("inner_2",),
config={
"configurable": {
"thread_id": "4",
"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": "4",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
),
StateSnapshot(
values={"my_key": "my value"},
@@ -8193,6 +8298,41 @@ def test_nested_graph_interrupts(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={
"my_key": "hi my value here and there",
"my_other_key": "hi my value here",
},
next=(),
config={
"configurable": {
"thread_id": "4",
"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": "4",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
),
StateSnapshot(
values={"my_key": "my value"},
@@ -8245,7 +8385,6 @@ def test_nested_graph_interrupts(
"my_key": "hi my value",
},
]
# interrupted after "inner"
assert list(app.get_state_history(config)) == [
StateSnapshot(
values={"my_key": "hi my value"},
@@ -8270,6 +8409,41 @@ def test_nested_graph_interrupts(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={
"my_key": "hi my value here",
"my_other_key": "hi my value",
},
next=("inner_2",),
config={
"configurable": {
"thread_id": "5",
"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": "5",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
),
StateSnapshot(
values={"my_key": "my value"},
@@ -8315,6 +8489,7 @@ def test_nested_graph_interrupts(
"my_key": "hi my value here and there",
},
]
# interrupted after "inner"
assert list(app.get_state_history(config)) == [
StateSnapshot(
values={"my_key": "hi my value here and there"},
@@ -8363,6 +8538,41 @@ def test_nested_graph_interrupts(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={
"my_key": "hi my value here and there",
"my_other_key": "hi my value here",
},
next=(),
config={
"configurable": {
"thread_id": "5",
"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": "5",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
),
StateSnapshot(
values={"my_key": "my value"},
@@ -8482,6 +8692,41 @@ def test_nested_graph_interrupts(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={
"my_key": "hi my value here and there",
"my_other_key": "hi my value here",
},
next=(),
config={
"configurable": {
"thread_id": "5",
"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": "5",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
),
StateSnapshot(
values={"my_key": "my value"},
@@ -8554,6 +8799,41 @@ def test_nested_graph_interrupts(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={
"my_key": "hi my value here",
"my_other_key": "hi my value",
},
next=("inner_2",),
config={
"configurable": {
"thread_id": "6",
"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": "6",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
),
StateSnapshot(
values={"my_key": "my value"},
@@ -8672,6 +8952,40 @@ def test_nested_graph_interrupts(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={
"my_key": "hi my value here",
"my_other_key": "hi my value",
},
next=("inner_2",),
config={
"configurable": {
"thread_id": "6",
"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": "6",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
),
},
),
StateSnapshot(
values={"my_key": "hi my value"},
@@ -8696,6 +9010,40 @@ def test_nested_graph_interrupts(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={
"my_key": "hi my value here",
"my_other_key": "hi my value",
},
next=("inner_2",),
config={
"configurable": {
"thread_id": "6",
"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": "6",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
),
},
),
StateSnapshot(
values={"my_key": "my value"},
@@ -8815,6 +9163,41 @@ def test_nested_graph_interrupts(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={
"my_key": "hi my value here",
"my_other_key": "hi my value",
},
next=("inner_2",),
config={
"configurable": {
"thread_id": "6",
"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": "6",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
),
StateSnapshot(
values={"my_key": "hi my value"},
@@ -8839,6 +9222,41 @@ def test_nested_graph_interrupts(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={
"my_key": "hi my value here and there",
"my_other_key": "hi my value here",
},
next=(),
config={
"configurable": {
"thread_id": "6",
"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": "6",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
),
StateSnapshot(
values={"my_key": "my value"},
@@ -9165,32 +9583,7 @@ def test_nested_graph_state(
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(
assert app.get_state(config) == StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={
@@ -9246,7 +9639,7 @@ def test_nested_graph_state(
)
},
)
assert list(app.get_state_history(config, include_subgraph_state=True)) == [
assert list(app.get_state_history(config)) == [
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
@@ -9349,7 +9742,7 @@ def test_nested_graph_state(
]
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(
assert app.get_state(config) == StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
next=(),
config={
@@ -9421,7 +9814,6 @@ def test_nested_graph_state(
"checkpoint_id": child_snapshot.config["configurable"]["checkpoint_id"],
}
},
include_subgraph_state=True,
) == StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
@@ -9448,7 +9840,7 @@ def test_nested_graph_state(
subgraph_state_snapshots={"inner": child_snapshot},
)
# test full history at the end
assert list(app.get_state_history(config, include_subgraph_state=True)) == [
assert list(app.get_state_history(config)) == [
StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
next=(),
@@ -9665,31 +10057,6 @@ def test_doubly_nested_graph_state(
config = {"configurable": {"thread_id": "1"}}
app.invoke({"my_key": "my value"}, config, debug=True)
assert app.get_state(config) == StateSnapshot(
values={"my_key": "hi my value"},
next=("child",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"parent_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=("child",),
config={
@@ -9763,7 +10130,7 @@ def test_doubly_nested_graph_state(
},
)
app.invoke(None, config, debug=True)
assert app.get_state(config, include_subgraph_state=True) == StateSnapshot(
assert app.get_state(config) == StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
next=(),
config={
@@ -9821,7 +10188,6 @@ def test_doubly_nested_graph_state(
# test getting child snapshot
child_snapshot = app.get_state(
{"configurable": {"thread_id": "1", "checkpoint_ns": "child"}},
include_subgraph_state=True,
)
assert child_snapshot == StateSnapshot(
values={"my_key": "hi my value here and there"},
@@ -9856,7 +10222,6 @@ def test_doubly_nested_graph_state(
"checkpoint_id": child_snapshot.config["configurable"]["checkpoint_id"],
}
},
include_subgraph_state=True,
) == StateSnapshot(
values={"my_key": "hi my value"},
next=("child",),
@@ -9935,7 +10300,7 @@ def test_send_to_nested_graphs(
"subjects": ["cats", "dogs"],
"jokes": [],
}
actual_snapshot = graph.get_state(config, include_subgraph_state=True)
actual_snapshot = graph.get_state(config)
subgraph_nodes = list(actual_snapshot.subgraph_state_snapshots.keys())
assert len(subgraph_nodes) == 2
for subgraph_node in subgraph_nodes:
@@ -9977,7 +10342,7 @@ def test_send_to_nested_graphs(
"jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"],
}
actual_snapshot = graph.get_state(config, include_subgraph_state=True)
actual_snapshot = graph.get_state(config)
expected_snapshot = StateSnapshot(
values={
"subjects": ["cats", "dogs"],
@@ -10013,7 +10378,7 @@ def test_send_to_nested_graphs(
assert actual_snapshot == expected_snapshot
# test full history
actual_history = list(graph.get_state_history(config, include_subgraph_state=True))
actual_history = list(graph.get_state_history(config))
# get subgraph node state for expected history
subgraph_state_snapshots = {
+427 -67
View File
@@ -6278,6 +6278,41 @@ async def test_nested_graph_interrupts(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={
"my_key": "hi my value here",
"my_other_key": "hi my value",
},
next=("inner_2",),
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"},
@@ -6395,6 +6430,41 @@ async def test_nested_graph_interrupts(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={
"my_key": "hi my value here and there",
"my_other_key": "hi my value here",
},
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"},
@@ -6571,6 +6641,41 @@ async def test_nested_graph_interrupts(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={
"my_key": "hi my value here",
"my_other_key": "hi my value",
},
next=("inner_2",),
config={
"configurable": {
"thread_id": "4",
"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": "4",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
),
StateSnapshot(
values={"my_key": "my value"},
@@ -6693,6 +6798,41 @@ async def test_nested_graph_interrupts(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={
"my_key": "hi my value here and there",
"my_other_key": "hi my value here",
},
next=(),
config={
"configurable": {
"thread_id": "4",
"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": "4",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
),
StateSnapshot(
values={"my_key": "my value"},
@@ -6772,6 +6912,41 @@ async def test_nested_graph_interrupts(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={
"my_key": "hi my value here",
"my_other_key": "hi my value",
},
next=("inner_2",),
config={
"configurable": {
"thread_id": "5",
"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": "5",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
),
StateSnapshot(
values={"my_key": "my value"},
@@ -6817,6 +6992,7 @@ async def test_nested_graph_interrupts(
"my_key": "hi my value here and there",
},
]
# interrupted after "inner"
assert [s async for s in app.aget_state_history(config)] == [
StateSnapshot(
values={"my_key": "hi my value here and there"},
@@ -6865,6 +7041,41 @@ async def test_nested_graph_interrupts(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={
"my_key": "hi my value here and there",
"my_other_key": "hi my value here",
},
next=(),
config={
"configurable": {
"thread_id": "5",
"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": "5",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
),
StateSnapshot(
values={"my_key": "my value"},
@@ -6984,6 +7195,41 @@ async def test_nested_graph_interrupts(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={
"my_key": "hi my value here and there",
"my_other_key": "hi my value here",
},
next=(),
config={
"configurable": {
"thread_id": "5",
"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": "5",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
),
StateSnapshot(
values={"my_key": "my value"},
@@ -7055,6 +7301,41 @@ async def test_nested_graph_interrupts(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={
"my_key": "hi my value here",
"my_other_key": "hi my value",
},
next=("inner_2",),
config={
"configurable": {
"thread_id": "6",
"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": "6",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
),
StateSnapshot(
values={"my_key": "my value"},
@@ -7174,6 +7455,40 @@ async def test_nested_graph_interrupts(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={
"my_key": "hi my value here",
"my_other_key": "hi my value",
},
next=("inner_2",),
config={
"configurable": {
"thread_id": "6",
"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": "6",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
),
},
),
StateSnapshot(
values={"my_key": "hi my value"},
@@ -7198,6 +7513,40 @@ async def test_nested_graph_interrupts(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={
"my_key": "hi my value here",
"my_other_key": "hi my value",
},
next=("inner_2",),
config={
"configurable": {
"thread_id": "6",
"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": "6",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
),
},
),
StateSnapshot(
values={"my_key": "my value"},
@@ -7317,6 +7666,41 @@ async def test_nested_graph_interrupts(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={
"my_key": "hi my value here",
"my_other_key": "hi my value",
},
next=("inner_2",),
config={
"configurable": {
"thread_id": "6",
"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": "6",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
),
StateSnapshot(
values={"my_key": "hi my value"},
@@ -7341,6 +7725,41 @@ async def test_nested_graph_interrupts(
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"inner": StateSnapshot(
values={
"my_key": "hi my value here and there",
"my_other_key": "hi my value here",
},
next=(),
config={
"configurable": {
"thread_id": "6",
"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": "6",
"checkpoint_ns": "inner",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
),
StateSnapshot(
values={"my_key": "my value"},
@@ -7672,32 +8091,7 @@ async def test_nested_graph_state(
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 await app.aget_state(config, include_subgraph_state=True) == StateSnapshot(
assert await app.aget_state(config) == StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={
@@ -7756,9 +8150,7 @@ async def test_nested_graph_state(
)
},
)
assert [
s async for s in app.aget_state_history(config, include_subgraph_state=True)
] == [
assert [s async for s in app.aget_state_history(config)] == [
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
@@ -7861,7 +8253,7 @@ async def test_nested_graph_state(
]
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(
assert await app.aget_state(config) == StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
next=(),
config={
@@ -7933,7 +8325,6 @@ async def test_nested_graph_state(
"checkpoint_id": child_snapshot.config["configurable"]["checkpoint_id"],
}
},
include_subgraph_state=True,
) == StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
@@ -7960,9 +8351,7 @@ async def test_nested_graph_state(
subgraph_state_snapshots={"inner": child_snapshot},
)
# test full history at the end
assert [
s async for s in app.aget_state_history(config, include_subgraph_state=True)
] == [
assert [s async for s in app.aget_state_history(config)] == [
StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
next=(),
@@ -8179,31 +8568,6 @@ async def test_doubly_nested_graph_state(
config = {"configurable": {"thread_id": "1"}}
await app.ainvoke({"my_key": "my value"}, config, debug=True)
assert await app.aget_state(config) == StateSnapshot(
values={"my_key": "hi my value"},
next=("child",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"parent_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 await app.aget_state(config, include_subgraph_state=True) == StateSnapshot(
values={"my_key": "hi my value"},
next=("child",),
config={
@@ -8277,7 +8641,7 @@ async def test_doubly_nested_graph_state(
},
)
await app.ainvoke(None, config, debug=True)
assert await app.aget_state(config, include_subgraph_state=True) == StateSnapshot(
assert await app.aget_state(config) == StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
next=(),
config={
@@ -8335,7 +8699,6 @@ async def test_doubly_nested_graph_state(
# test getting child snapshot
child_snapshot = await app.aget_state(
{"configurable": {"thread_id": "1", "checkpoint_ns": "child"}},
include_subgraph_state=True,
)
assert child_snapshot == StateSnapshot(
values={"my_key": "hi my value here and there"},
@@ -8370,7 +8733,6 @@ async def test_doubly_nested_graph_state(
"checkpoint_id": child_snapshot.config["configurable"]["checkpoint_id"],
}
},
include_subgraph_state=True,
) == StateSnapshot(
values={"my_key": "hi my value"},
next=("child",),
@@ -8449,7 +8811,7 @@ async def test_send_to_nested_graphs(
"subjects": ["cats", "dogs"],
"jokes": [],
}
actual_snapshot = await graph.aget_state(config, include_subgraph_state=True)
actual_snapshot = await graph.aget_state(config)
subgraph_nodes = list(actual_snapshot.subgraph_state_snapshots.keys())
assert len(subgraph_nodes) == 2
for subgraph_node in subgraph_nodes:
@@ -8490,7 +8852,7 @@ async def test_send_to_nested_graphs(
"jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"],
}
actual_snapshot = await graph.aget_state(config, include_subgraph_state=True)
actual_snapshot = await graph.aget_state(config)
expected_snapshot = StateSnapshot(
values={
"subjects": ["cats", "dogs"],
@@ -8526,9 +8888,7 @@ async def test_send_to_nested_graphs(
assert actual_snapshot == expected_snapshot
# test full history
actual_history = [
c async for c in graph.aget_state_history(config, include_subgraph_state=True)
]
actual_history = [c async for c in graph.aget_state_history(config)]
# get subgraph node state for expected history
subgraph_state_snapshots = {
subgraph_node: await graph.aget_state(