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

Merge pull request #1530 from langchain-ai/nc/29aug/get-state-subgraph

Browse Source 
Enable passing subgraph checkpoint config to get_state/aget_state
This commit is contained in:
Nuno Campos
2024-08-29 11:59:02 -07:00
committed by GitHub
parent 01f14d3166 9d68aac76f
commit a627a30f4c
3 changed files with 208 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/__init__.py
+48 -4
View File
@@ -370,9 +370,9 @@ class Pregel(
)
async def aget_subgraphs(
self, recursive: bool = False
self, recurse: bool = False
) -> AsyncIterator[tuple[str, Pregel]]:
for name, node in self.get_subgraphs(recurse=recursive):
for name, node in self.get_subgraphs(recurse=recurse):
yield name, node
def _prepare_state_snapshot(
@@ -482,7 +482,7 @@ class Pregel(
for_execution=False,
)
# get the subgraphs
subgraphs = dict(self.get_subgraphs())
subgraphs = {n: g async for n, g in self.aget_subgraphs()}
parent_ns = saved.config["configurable"].get("checkpoint_ns", "")
task_states: dict[str, Union[RunnableConfig, StateSnapshot]] = {}
for task in next_tasks:
@@ -534,6 +534,28 @@ class Pregel(
if not checkpointer:
raise ValueError("No checkpointer set")
if (
checkpoint_ns := config["configurable"].get("checkpoint_ns", "")
) and CONFIG_KEY_CHECKPOINTER not in config["configurable"]:
# remove task_ids from checkpoint_ns
recast_checkpoint_ns = NS_SEP.join(
part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
)
# find the subgraph with the matching name
for name, pregel in self.get_subgraphs(recurse=True):
if name == recast_checkpoint_ns:
return pregel.get_state(
{
"configurable": {
**config["configurable"],
CONFIG_KEY_CHECKPOINTER: checkpointer,
}
},
subgraphs=subgraphs,
)
else:
raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
config = merge_configs(self.config, config) if self.config else config
saved = checkpointer.get_tuple(config)
return self._prepare_state_snapshot(
@@ -550,6 +572,28 @@ class Pregel(
if not checkpointer:
raise ValueError("No checkpointer set")
if (
checkpoint_ns := config["configurable"].get("checkpoint_ns", "")
) and CONFIG_KEY_CHECKPOINTER not in config["configurable"]:
# remove task_ids from checkpoint_ns
recast_checkpoint_ns = NS_SEP.join(
part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
)
# find the subgraph with the matching name
async for name, pregel in self.aget_subgraphs(recurse=True):
if name == recast_checkpoint_ns:
return await pregel.aget_state(
{
"configurable": {
**config["configurable"],
CONFIG_KEY_CHECKPOINTER: checkpointer,
}
},
subgraphs=subgraphs,
)
else:
raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
config = merge_configs(self.config, config) if self.config else config
saved = await checkpointer.aget_tuple(config)
return await self._aprepare_state_snapshot(
@@ -630,7 +674,7 @@ class Pregel(
part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP)
)
# find the subgraph with the matching name
for name, pregel in self.get_subgraphs(recurse=True):
async for name, pregel in self.aget_subgraphs(recurse=True):
if name == recast_checkpoint_ns:
async for state in pregel.aget_state_history(
{
libs/langgraph/tests/test_pregel.py
+80 -1
View File
@@ -10899,7 +10899,8 @@ def test_doubly_nested_graph_state(
config = {"configurable": {"thread_id": "1"}}
app.invoke({"my_key": "my value"}, config, debug=True)
# get state without subgraphs
assert app.get_state(config) == StateSnapshot(
outer_state = app.get_state(config)
assert outer_state == StateSnapshot(
values={"my_key": "hi my value"},
tasks=(
PregelTask(
@@ -10936,6 +10937,84 @@ def test_doubly_nested_graph_state(
}
},
)
child_state = app.get_state(outer_state.tasks[0].state)
assert (
child_state.tasks[0]
== StateSnapshot(
values={"my_key": "hi my value"},
tasks=(
PregelTask(
AnyStr(),
"child_1",
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr(),
}
},
),
),
next=("child_1",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("child:"),
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {"": AnyStr()},
"source": "loop",
"writes": None,
"step": 0,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("child:"),
"checkpoint_id": AnyStr(),
}
},
).tasks[0]
)
grandchild_state = app.get_state(child_state.tasks[0].state)
assert grandchild_state == StateSnapshot(
values={"my_key": "hi my value here"},
tasks=(
PregelTask(
AnyStr(),
"grandchild_2",
),
),
next=("grandchild_2",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr(),
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": AnyDict(
{
"": AnyStr(),
AnyStr("child:"): AnyStr(),
}
),
"source": "loop",
"writes": {"grandchild_1": {"my_key": "hi my value here"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr(),
"checkpoint_id": AnyStr(),
}
},
)
# get state with subgraphs
assert app.get_state(config, subgraphs=True) == StateSnapshot(
values={"my_key": "hi my value"},
libs/langgraph/tests/test_pregel_async.py
+80 -1
View File
@@ -9340,7 +9340,8 @@ async def test_doubly_nested_graph_state(
config = {"configurable": {"thread_id": "1"}}
await app.ainvoke({"my_key": "my value"}, config, debug=True)
# get state without subgraphs
assert await app.aget_state(config) == StateSnapshot(
outer_state = await app.aget_state(config)
assert outer_state == StateSnapshot(
values={"my_key": "hi my value"},
tasks=(
PregelTask(
@@ -9377,6 +9378,84 @@ async def test_doubly_nested_graph_state(
}
},
)
child_state = await app.aget_state(outer_state.tasks[0].state)
assert (
child_state.tasks[0]
== StateSnapshot(
values={"my_key": "hi my value"},
tasks=(
PregelTask(
AnyStr(),
"child_1",
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr(),
}
},
),
),
next=("child_1",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("child:"),
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {"": AnyStr()},
"source": "loop",
"writes": None,
"step": 0,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("child:"),
"checkpoint_id": AnyStr(),
}
},
).tasks[0]
)
grandchild_state = await app.aget_state(child_state.tasks[0].state)
assert grandchild_state == StateSnapshot(
values={"my_key": "hi my value here"},
tasks=(
PregelTask(
AnyStr(),
"grandchild_2",
),
),
next=("grandchild_2",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr(),
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": AnyDict(
{
"": AnyStr(),
AnyStr("child:"): AnyStr(),
}
),
"source": "loop",
"writes": {"grandchild_1": {"my_key": "hi my value here"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr(),
"checkpoint_id": AnyStr(),
}
},
)
# get state with subgraphs
assert await app.aget_state(config, subgraphs=True) == StateSnapshot(
values={"my_key": "hi my value"},