mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-10-01 14:05:10 +02:00
Replay a nested checkpoint
- note this isn't very useful until update_state is added
This commit is contained in:
@@ -59,11 +59,13 @@ from langgraph.checkpoint.base import (
|
||||
empty_checkpoint,
|
||||
)
|
||||
from langgraph.constants import (
|
||||
CONFIG_KEY_CHECKPOINT_MAP,
|
||||
CONFIG_KEY_CHECKPOINTER,
|
||||
CONFIG_KEY_READ,
|
||||
CONFIG_KEY_RESUMING,
|
||||
CONFIG_KEY_SEND,
|
||||
CONFIG_KEY_STREAM,
|
||||
CONFIG_KEY_TASK_ID,
|
||||
ERROR,
|
||||
INTERRUPT,
|
||||
NS_END,
|
||||
@@ -77,6 +79,7 @@ from langgraph.pregel.algo import (
|
||||
local_write,
|
||||
prepare_next_tasks,
|
||||
)
|
||||
from langgraph.pregel.config import patch_configurable
|
||||
from langgraph.pregel.debug import (
|
||||
print_step_checkpoint,
|
||||
print_step_tasks,
|
||||
@@ -442,7 +445,19 @@ class Pregel(
|
||||
return StateSnapshot(
|
||||
read_channels(channels, self.stream_channels_asis),
|
||||
tuple(t.name for t in next_tasks),
|
||||
saved.config,
|
||||
patch_configurable(
|
||||
saved.config,
|
||||
{
|
||||
CONFIG_KEY_CHECKPOINT_MAP: {
|
||||
**saved.metadata["parents"],
|
||||
saved.config["configurable"][
|
||||
"checkpoint_ns"
|
||||
]: saved.checkpoint["id"],
|
||||
}
|
||||
},
|
||||
)
|
||||
if saved.metadata.get("parents")
|
||||
else saved.config,
|
||||
saved.metadata,
|
||||
saved.checkpoint["ts"],
|
||||
saved.parent_config,
|
||||
@@ -518,7 +533,19 @@ class Pregel(
|
||||
return StateSnapshot(
|
||||
read_channels(channels, self.stream_channels_asis),
|
||||
tuple(t.name for t in next_tasks),
|
||||
saved.config,
|
||||
patch_configurable(
|
||||
saved.config,
|
||||
{
|
||||
CONFIG_KEY_CHECKPOINT_MAP: {
|
||||
**saved.metadata["parents"],
|
||||
saved.config["configurable"][
|
||||
"checkpoint_ns"
|
||||
]: saved.checkpoint["id"],
|
||||
}
|
||||
},
|
||||
)
|
||||
if saved.metadata.get("parents")
|
||||
else saved.config,
|
||||
saved.metadata,
|
||||
saved.checkpoint["ts"],
|
||||
saved.parent_config,
|
||||
@@ -546,12 +573,9 @@ class Pregel(
|
||||
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,
|
||||
}
|
||||
},
|
||||
patch_configurable(
|
||||
config, {CONFIG_KEY_CHECKPOINTER: checkpointer}
|
||||
),
|
||||
subgraphs=subgraphs,
|
||||
)
|
||||
else:
|
||||
@@ -584,12 +608,9 @@ class Pregel(
|
||||
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,
|
||||
}
|
||||
},
|
||||
patch_configurable(
|
||||
config, {CONFIG_KEY_CHECKPOINTER: checkpointer}
|
||||
),
|
||||
subgraphs=subgraphs,
|
||||
)
|
||||
else:
|
||||
@@ -627,12 +648,9 @@ class Pregel(
|
||||
for name, pregel in self.get_subgraphs(recurse=True):
|
||||
if name == recast_checkpoint_ns:
|
||||
yield from pregel.get_state_history(
|
||||
{
|
||||
"configurable": {
|
||||
**config["configurable"],
|
||||
CONFIG_KEY_CHECKPOINTER: checkpointer,
|
||||
}
|
||||
},
|
||||
patch_configurable(
|
||||
config, {CONFIG_KEY_CHECKPOINTER: checkpointer}
|
||||
),
|
||||
filter=filter,
|
||||
before=before,
|
||||
limit=limit,
|
||||
@@ -678,12 +696,9 @@ class Pregel(
|
||||
async for name, pregel in self.aget_subgraphs(recurse=True):
|
||||
if name == recast_checkpoint_ns:
|
||||
async for state in pregel.aget_state_history(
|
||||
{
|
||||
"configurable": {
|
||||
**config["configurable"],
|
||||
CONFIG_KEY_CHECKPOINTER: checkpointer,
|
||||
}
|
||||
},
|
||||
patch_configurable(
|
||||
config, {CONFIG_KEY_CHECKPOINTER: checkpointer}
|
||||
),
|
||||
filter=filter,
|
||||
before=before,
|
||||
limit=limit,
|
||||
@@ -717,9 +732,32 @@ class Pregel(
|
||||
node `as_node`. If `as_node` is not provided, it will be set to the last node
|
||||
that updated the state, if not ambiguous.
|
||||
"""
|
||||
if not self.checkpointer:
|
||||
checkpointer: Optional[BaseCheckpointSaver] = config["configurable"].get(
|
||||
CONFIG_KEY_CHECKPOINTER, self.checkpointer
|
||||
)
|
||||
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.update_state(
|
||||
patch_configurable(
|
||||
config, {CONFIG_KEY_CHECKPOINTER: checkpointer}
|
||||
),
|
||||
values,
|
||||
as_node,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
|
||||
|
||||
# get last checkpoint
|
||||
config = merge_configs(self.config, config) if self.config else config
|
||||
saved = self.checkpointer.get_tuple(config)
|
||||
@@ -729,21 +767,12 @@ class Pregel(
|
||||
)
|
||||
step = saved.metadata.get("step", -1) if saved else -1
|
||||
# merge configurable fields with previous checkpoint config
|
||||
checkpoint_config = {
|
||||
**config,
|
||||
"configurable": {
|
||||
**config["configurable"],
|
||||
# TODO: add proper support for updating nested subgraph state
|
||||
"checkpoint_ns": "",
|
||||
},
|
||||
}
|
||||
checkpoint_config = patch_configurable(
|
||||
config,
|
||||
{"checkpoint_ns": config["configurable"].get("checkpoint_ns", "")},
|
||||
)
|
||||
if saved:
|
||||
checkpoint_config = {
|
||||
"configurable": {
|
||||
**config.get("configurable", {}),
|
||||
**saved.config["configurable"],
|
||||
}
|
||||
}
|
||||
checkpoint_config = patch_configurable(config, saved.config["configurable"])
|
||||
# find last node that updated the state, if not provided
|
||||
if values is None and as_node is None:
|
||||
return self.checkpointer.put(
|
||||
@@ -798,6 +827,7 @@ class Pregel(
|
||||
None,
|
||||
[INTERRUPT],
|
||||
None,
|
||||
None,
|
||||
str(uuid5(UUID(checkpoint["id"]), INTERRUPT)),
|
||||
)
|
||||
# execute task
|
||||
@@ -935,6 +965,7 @@ class Pregel(
|
||||
None,
|
||||
[INTERRUPT],
|
||||
None,
|
||||
None,
|
||||
str(uuid5(UUID(checkpoint["id"]), INTERRUPT)),
|
||||
)
|
||||
# execute task
|
||||
@@ -1016,7 +1047,7 @@ class Pregel(
|
||||
stream_mode = stream_mode if stream_mode is not None else self.stream_mode
|
||||
if not isinstance(stream_mode, list):
|
||||
stream_mode = [stream_mode]
|
||||
if CONFIG_KEY_READ in config.get("configurable", {}):
|
||||
if CONFIG_KEY_TASK_ID in config.get("configurable", {}):
|
||||
# if being called as a node in another graph, always use values mode
|
||||
stream_mode = ["values"]
|
||||
if CONFIG_KEY_CHECKPOINTER in config.get("configurable", {}):
|
||||
|
||||
@@ -50,7 +50,7 @@ from langgraph.pregel.io import read_channel, read_channels
|
||||
from langgraph.pregel.log import logger
|
||||
from langgraph.pregel.manager import ChannelsManager
|
||||
from langgraph.pregel.read import PregelNode
|
||||
from langgraph.pregel.types import All, PregelExecutableTask, PregelTask
|
||||
from langgraph.pregel.types import EXACT_MATCH, All, PregelExecutableTask, PregelTask
|
||||
|
||||
|
||||
class WritesProtocol(Protocol):
|
||||
@@ -303,6 +303,7 @@ def prepare_next_tasks(
|
||||
if node := proc.get_node():
|
||||
managed.replace_runtime_placeholders(step, packet.arg)
|
||||
writes = deque()
|
||||
task_checkpoint_ns = f"{checkpoint_ns}:{task_id}"
|
||||
tasks.append(
|
||||
PregelExecutableTask(
|
||||
packet.node,
|
||||
@@ -351,11 +352,15 @@ def prepare_next_tasks(
|
||||
},
|
||||
CONFIG_KEY_RESUMING: is_resuming,
|
||||
"checkpoint_id": None,
|
||||
"checkpoint_ns": f"{checkpoint_ns}:{task_id}",
|
||||
"checkpoint_ns": task_checkpoint_ns,
|
||||
},
|
||||
),
|
||||
triggers,
|
||||
proc.retry_policy,
|
||||
None
|
||||
if task_checkpoint_ns
|
||||
in configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {})
|
||||
else EXACT_MATCH,
|
||||
task_id,
|
||||
)
|
||||
)
|
||||
@@ -406,6 +411,7 @@ def prepare_next_tasks(
|
||||
if for_execution:
|
||||
if node := proc.get_node():
|
||||
writes = deque()
|
||||
task_checkpoint_ns = f"{checkpoint_ns}:{task_id}"
|
||||
tasks.append(
|
||||
PregelExecutableTask(
|
||||
name,
|
||||
@@ -455,11 +461,15 @@ def prepare_next_tasks(
|
||||
parent_ns: checkpoint["id"],
|
||||
},
|
||||
CONFIG_KEY_RESUMING: is_resuming,
|
||||
"checkpoint_ns": f"{checkpoint_ns}:{task_id}",
|
||||
"checkpoint_ns": task_checkpoint_ns,
|
||||
},
|
||||
),
|
||||
triggers,
|
||||
proc.retry_policy,
|
||||
None
|
||||
if task_checkpoint_ns
|
||||
in configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {})
|
||||
else EXACT_MATCH,
|
||||
task_id,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
from typing import Any, Optional
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
|
||||
|
||||
def patch_configurable(
|
||||
config: Optional[RunnableConfig], patch: dict[str, Any]
|
||||
) -> RunnableConfig:
|
||||
if config is None:
|
||||
return {"configurable": patch}
|
||||
else:
|
||||
return {**config, "configurable": {**config["configurable"], **patch}}
|
||||
@@ -78,11 +78,11 @@ def map_debug_tasks(
|
||||
step: int, tasks: list[PregelExecutableTask]
|
||||
) -> Iterator[DebugOutputTask]:
|
||||
ts = datetime.now(timezone.utc).isoformat()
|
||||
for name, input, _, _, config, triggers, _, _ in tasks:
|
||||
if config is not None and TAG_HIDDEN in config.get("tags", []):
|
||||
for task in tasks:
|
||||
if task.config is not None and TAG_HIDDEN in task.config.get("tags", []):
|
||||
continue
|
||||
|
||||
metadata = config["metadata"].copy()
|
||||
metadata = task.config["metadata"].copy()
|
||||
metadata.pop("checkpoint_id", None)
|
||||
|
||||
yield {
|
||||
@@ -90,10 +90,12 @@ def map_debug_tasks(
|
||||
"timestamp": ts,
|
||||
"step": step,
|
||||
"payload": {
|
||||
"id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step, metadata)))),
|
||||
"name": name,
|
||||
"input": input,
|
||||
"triggers": triggers,
|
||||
"id": str(
|
||||
uuid5(TASK_NAMESPACE, json.dumps((task.name, step, metadata)))
|
||||
),
|
||||
"name": task.name,
|
||||
"input": task.input,
|
||||
"triggers": task.triggers,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -107,11 +109,11 @@ def map_debug_task_results(
|
||||
[stream_keys] if isinstance(stream_keys, str) else stream_keys
|
||||
)
|
||||
ts = datetime.now(timezone.utc).isoformat()
|
||||
for (name, _, _, _, config, _, _, _), writes in tasks:
|
||||
if config is not None and TAG_HIDDEN in config.get("tags", []):
|
||||
for task, writes in tasks:
|
||||
if task.config is not None and TAG_HIDDEN in task.config.get("tags", []):
|
||||
continue
|
||||
|
||||
metadata = config["metadata"].copy()
|
||||
metadata = task.config["metadata"].copy()
|
||||
metadata.pop("checkpoint_id", None)
|
||||
# TODO: make task IDs deterministic in tests and reuse task IDs for payload ID
|
||||
|
||||
@@ -120,8 +122,10 @@ def map_debug_task_results(
|
||||
"timestamp": ts,
|
||||
"step": step,
|
||||
"payload": {
|
||||
"id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step, metadata)))),
|
||||
"name": name,
|
||||
"id": str(
|
||||
uuid5(TASK_NAMESPACE, json.dumps((task.name, step, metadata)))
|
||||
),
|
||||
"name": task.name,
|
||||
"error": next((w[1] for w in writes if w[0] == ERROR), None),
|
||||
"result": [w for w in writes if w[0] in stream_channels_list],
|
||||
"interrupts": [asdict(w[1]) for w in writes if w[0] == INTERRUPT],
|
||||
|
||||
@@ -40,9 +40,9 @@ from langgraph.checkpoint.base import (
|
||||
)
|
||||
from langgraph.constants import (
|
||||
CONFIG_KEY_CHECKPOINT_MAP,
|
||||
CONFIG_KEY_READ,
|
||||
CONFIG_KEY_RESUMING,
|
||||
CONFIG_KEY_STREAM,
|
||||
CONFIG_KEY_TASK_ID,
|
||||
ERROR,
|
||||
INPUT,
|
||||
INTERRUPT,
|
||||
@@ -60,6 +60,7 @@ from langgraph.pregel.algo import (
|
||||
prepare_next_tasks,
|
||||
should_interrupt,
|
||||
)
|
||||
from langgraph.pregel.config import patch_configurable
|
||||
from langgraph.pregel.debug import (
|
||||
map_debug_checkpoint,
|
||||
map_debug_task_results,
|
||||
@@ -178,11 +179,30 @@ class PregelLoop:
|
||||
self.specs = specs
|
||||
self.output_keys = output_keys
|
||||
self.stream_keys = stream_keys
|
||||
self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {})
|
||||
self.is_nested = CONFIG_KEY_TASK_ID in self.config.get("configurable", {})
|
||||
if CONFIG_KEY_STREAM in config["configurable"]:
|
||||
self.stream = DuplexStream(
|
||||
self.stream, config["configurable"][CONFIG_KEY_STREAM]
|
||||
)
|
||||
if not self.is_nested and config["configurable"].get("checkpoint_ns"):
|
||||
self.config = patch_configurable(
|
||||
config, {"checkpoint_ns": "", "checkpoint_id": None}
|
||||
)
|
||||
if (
|
||||
CONFIG_KEY_CHECKPOINT_MAP in self.config["configurable"]
|
||||
and self.config["configurable"].get("checkpoint_ns")
|
||||
in self.config["configurable"][CONFIG_KEY_CHECKPOINT_MAP]
|
||||
):
|
||||
self.checkpoint_config = patch_configurable(
|
||||
self.config,
|
||||
{
|
||||
"checkpoint_id": config["configurable"][CONFIG_KEY_CHECKPOINT_MAP][
|
||||
self.config["configurable"]["checkpoint_ns"]
|
||||
]
|
||||
},
|
||||
)
|
||||
else:
|
||||
self.checkpoint_config = config
|
||||
|
||||
def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None:
|
||||
"""Put writes for a task, to be read by the next tick."""
|
||||
@@ -324,7 +344,14 @@ class PregelLoop:
|
||||
for tid, k, v in self.checkpoint_pending_writes:
|
||||
if k in (ERROR, INTERRUPT):
|
||||
continue
|
||||
if task := next((t for t in self.tasks if t.id == tid), None):
|
||||
if task := next(
|
||||
(
|
||||
t
|
||||
for t in self.tasks
|
||||
if t.id == tid and t.cache_policy is not None
|
||||
),
|
||||
None,
|
||||
):
|
||||
task.writes.append((k, v))
|
||||
# print output for any tasks we applied previous writes to
|
||||
for task in self.tasks:
|
||||
@@ -524,7 +551,9 @@ class SyncPregelLoop(PregelLoop, ContextManager):
|
||||
|
||||
def __enter__(self) -> Self:
|
||||
saved = (
|
||||
self.checkpointer.get_tuple(self.config) if self.checkpointer else None
|
||||
self.checkpointer.get_tuple(self.checkpoint_config)
|
||||
if self.checkpointer
|
||||
else None
|
||||
) or CheckpointTuple(self.config, empty_checkpoint(), {"step": -2}, None, [])
|
||||
self.checkpoint_config = {
|
||||
**self.config,
|
||||
@@ -616,7 +645,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
|
||||
|
||||
async def __aenter__(self) -> Self:
|
||||
saved = (
|
||||
await self.checkpointer.aget_tuple(self.config)
|
||||
await self.checkpointer.aget_tuple(self.checkpoint_config)
|
||||
if self.checkpointer
|
||||
else None
|
||||
) or CheckpointTuple(self.config, empty_checkpoint(), {"step": -2}, None, [])
|
||||
|
||||
@@ -57,6 +57,15 @@ class RetryPolicy(NamedTuple):
|
||||
"""List of exception classes that should trigger a retry, or a callable that returns True for exceptions that should trigger a retry."""
|
||||
|
||||
|
||||
class CachePolicy(NamedTuple):
|
||||
"""Configuration for caching nodes."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
EXACT_MATCH = CachePolicy()
|
||||
|
||||
|
||||
class PregelTask(NamedTuple):
|
||||
id: str
|
||||
name: str
|
||||
@@ -73,6 +82,7 @@ class PregelExecutableTask(NamedTuple):
|
||||
config: RunnableConfig
|
||||
triggers: list[str]
|
||||
retry_policy: Optional[RetryPolicy]
|
||||
cache_policy: Optional[CachePolicy]
|
||||
id: str
|
||||
|
||||
|
||||
|
||||
@@ -1,16 +1,21 @@
|
||||
import re
|
||||
from typing import Any, Sequence
|
||||
|
||||
|
||||
class AnyStr(str):
|
||||
def __init__(self, prefix: str = "") -> None:
|
||||
def __init__(self, prefix: str | re.Pattern = "") -> None:
|
||||
super().__init__()
|
||||
self.prefix = prefix
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
return isinstance(other, str) and other.startswith(self.prefix)
|
||||
return isinstance(other, str) and (
|
||||
other.startswith(self.prefix)
|
||||
if isinstance(self.prefix, str)
|
||||
else self.prefix.match(other)
|
||||
)
|
||||
|
||||
def __hash__(self) -> int:
|
||||
return hash(str(self))
|
||||
return hash((str(self), self.prefix))
|
||||
|
||||
|
||||
class AnyDict(dict):
|
||||
|
||||
+87
-1815
File diff suppressed because it is too large
Load Diff
@@ -1,6 +1,7 @@
|
||||
import asyncio
|
||||
import json
|
||||
import operator
|
||||
import re
|
||||
from collections import Counter
|
||||
from contextlib import asynccontextmanager, contextmanager
|
||||
from typing import (
|
||||
@@ -8892,6 +8893,9 @@ async def test_nested_graph_state(
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": AnyStr("inner:"),
|
||||
"checkpoint_id": AnyStr(),
|
||||
"checkpoint_map": AnyDict(
|
||||
{"": AnyStr(), AnyStr("child:"): AnyStr()}
|
||||
),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
@@ -9039,6 +9043,9 @@ async def test_nested_graph_state(
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": AnyStr("inner:"),
|
||||
"checkpoint_id": AnyStr(),
|
||||
"checkpoint_map": AnyDict(
|
||||
{"": AnyStr(), AnyStr("inner:"): AnyStr()}
|
||||
),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
@@ -9070,6 +9077,9 @@ async def test_nested_graph_state(
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": AnyStr("inner:"),
|
||||
"checkpoint_id": AnyStr(),
|
||||
"checkpoint_map": AnyDict(
|
||||
{"": AnyStr(), AnyStr("inner:"): AnyStr()}
|
||||
),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
@@ -9096,6 +9106,9 @@ async def test_nested_graph_state(
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": AnyStr("inner:"),
|
||||
"checkpoint_id": AnyStr(),
|
||||
"checkpoint_map": AnyDict(
|
||||
{"": AnyStr(), AnyStr("inner:"): AnyStr()}
|
||||
),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
@@ -9413,6 +9426,9 @@ async def test_doubly_nested_graph_state(
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": AnyStr("child:"),
|
||||
"checkpoint_id": AnyStr(),
|
||||
"checkpoint_map": AnyDict(
|
||||
{"": AnyStr(), AnyStr("child:"): AnyStr()}
|
||||
),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
@@ -9446,6 +9462,13 @@ async def test_doubly_nested_graph_state(
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": AnyStr(),
|
||||
"checkpoint_id": AnyStr(),
|
||||
"checkpoint_map": AnyDict(
|
||||
{
|
||||
"": AnyStr(),
|
||||
AnyStr("child:"): AnyStr(),
|
||||
AnyStr(re.compile(r"child:.+|child1:")): AnyStr(),
|
||||
}
|
||||
),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
@@ -9495,6 +9518,15 @@ async def test_doubly_nested_graph_state(
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": AnyStr(),
|
||||
"checkpoint_id": AnyStr(),
|
||||
"checkpoint_map": AnyDict(
|
||||
{
|
||||
"": AnyStr(),
|
||||
AnyStr("child:"): AnyStr(),
|
||||
AnyStr(
|
||||
re.compile(r"child:.+|child1:")
|
||||
): AnyStr(),
|
||||
}
|
||||
),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
@@ -9527,6 +9559,9 @@ async def test_doubly_nested_graph_state(
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": AnyStr("child:"),
|
||||
"checkpoint_id": AnyStr(),
|
||||
"checkpoint_map": AnyDict(
|
||||
{"": AnyStr(), AnyStr("child:"): AnyStr()}
|
||||
),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
@@ -9614,145 +9649,150 @@ async def test_doubly_nested_graph_state(
|
||||
)
|
||||
# get outer graph history
|
||||
outer_history = [c async for c in app.aget_state_history(config)]
|
||||
assert outer_history == [
|
||||
StateSnapshot(
|
||||
values={"my_key": "hi my value here and there and back again"},
|
||||
tasks=(),
|
||||
next=(),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "loop",
|
||||
"writes": {
|
||||
"parent_2": {"my_key": "hi my value here and there and back again"}
|
||||
assert (
|
||||
outer_history[0]
|
||||
== [
|
||||
StateSnapshot(
|
||||
values={"my_key": "hi my value here and there and back again"},
|
||||
tasks=(),
|
||||
next=(),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
"step": 3,
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
parent_config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
),
|
||||
StateSnapshot(
|
||||
values={"my_key": "hi my value here and there"},
|
||||
next=("parent_2",),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "loop",
|
||||
"writes": {"child": {"my_key": "hi my value here and there"}},
|
||||
"step": 2,
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
parent_config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
tasks=(PregelTask(id=AnyStr(), name="parent_2"),),
|
||||
),
|
||||
StateSnapshot(
|
||||
values={"my_key": "hi my value"},
|
||||
tasks=(
|
||||
PregelTask(
|
||||
AnyStr(),
|
||||
"child",
|
||||
state={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": AnyStr("child"),
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "loop",
|
||||
"writes": {
|
||||
"parent_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(),
|
||||
}
|
||||
},
|
||||
),
|
||||
next=("child",),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
"parents": {},
|
||||
"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(),
|
||||
}
|
||||
},
|
||||
),
|
||||
StateSnapshot(
|
||||
values={"my_key": "my value"},
|
||||
next=("parent_1",),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "loop",
|
||||
"writes": None,
|
||||
"step": 0,
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
parent_config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
tasks=(PregelTask(id=AnyStr(), name="parent_1"),),
|
||||
),
|
||||
StateSnapshot(
|
||||
values={},
|
||||
next=("__start__",),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "input",
|
||||
"writes": {"my_key": "my value"},
|
||||
"step": -1,
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
parent_config=None,
|
||||
tasks=(PregelTask(id=AnyStr(), name="__start__"),),
|
||||
),
|
||||
]
|
||||
StateSnapshot(
|
||||
values={"my_key": "hi my value here and there"},
|
||||
next=("parent_2",),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "loop",
|
||||
"writes": {"child": {"my_key": "hi my value here and there"}},
|
||||
"step": 2,
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
parent_config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
tasks=(PregelTask(id=AnyStr(), name="parent_2"),),
|
||||
),
|
||||
StateSnapshot(
|
||||
values={"my_key": "hi my value"},
|
||||
tasks=(
|
||||
PregelTask(
|
||||
AnyStr(),
|
||||
"child",
|
||||
state={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": AnyStr("child"),
|
||||
}
|
||||
},
|
||||
),
|
||||
),
|
||||
next=("child",),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
"parents": {},
|
||||
"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(),
|
||||
}
|
||||
},
|
||||
),
|
||||
StateSnapshot(
|
||||
values={"my_key": "my value"},
|
||||
next=("parent_1",),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "loop",
|
||||
"writes": None,
|
||||
"step": 0,
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
parent_config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
tasks=(PregelTask(id=AnyStr(), name="parent_1"),),
|
||||
),
|
||||
StateSnapshot(
|
||||
values={},
|
||||
next=("__start__",),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "input",
|
||||
"writes": {"my_key": "my value"},
|
||||
"step": -1,
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
parent_config=None,
|
||||
tasks=(PregelTask(id=AnyStr(), name="__start__"),),
|
||||
),
|
||||
][0]
|
||||
)
|
||||
# get child graph history
|
||||
child_history = [
|
||||
c async for c in app.aget_state_history(outer_history[2].tasks[0].state)
|
||||
@@ -9766,6 +9806,9 @@ async def test_doubly_nested_graph_state(
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": AnyStr("child:"),
|
||||
"checkpoint_id": AnyStr(),
|
||||
"checkpoint_map": AnyDict(
|
||||
{"": AnyStr(), AnyStr("child:"): AnyStr()}
|
||||
),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
@@ -9792,6 +9835,9 @@ async def test_doubly_nested_graph_state(
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": AnyStr("child:"),
|
||||
"checkpoint_id": AnyStr(),
|
||||
"checkpoint_map": AnyDict(
|
||||
{"": AnyStr(), AnyStr("child:"): AnyStr()}
|
||||
),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
@@ -9829,6 +9875,9 @@ async def test_doubly_nested_graph_state(
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": AnyStr("child:"),
|
||||
"checkpoint_id": AnyStr(),
|
||||
"checkpoint_map": AnyDict(
|
||||
{"": AnyStr(), AnyStr("child:"): AnyStr()}
|
||||
),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
@@ -9855,6 +9904,13 @@ async def test_doubly_nested_graph_state(
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": AnyStr(),
|
||||
"checkpoint_id": AnyStr(),
|
||||
"checkpoint_map": AnyDict(
|
||||
{
|
||||
"": AnyStr(),
|
||||
AnyStr("child:"): AnyStr(),
|
||||
AnyStr(re.compile(r"child:.+|child1:")): AnyStr(),
|
||||
}
|
||||
),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
@@ -9886,6 +9942,13 @@ async def test_doubly_nested_graph_state(
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": AnyStr(),
|
||||
"checkpoint_id": AnyStr(),
|
||||
"checkpoint_map": AnyDict(
|
||||
{
|
||||
"": AnyStr(),
|
||||
AnyStr("child:"): AnyStr(),
|
||||
AnyStr(re.compile(r"child:.+|child1:")): AnyStr(),
|
||||
}
|
||||
),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
@@ -9917,6 +9980,13 @@ async def test_doubly_nested_graph_state(
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": AnyStr(),
|
||||
"checkpoint_id": AnyStr(),
|
||||
"checkpoint_map": AnyDict(
|
||||
{
|
||||
"": AnyStr(),
|
||||
AnyStr("child:"): AnyStr(),
|
||||
AnyStr(re.compile(r"child:.+|child1:")): AnyStr(),
|
||||
}
|
||||
),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
@@ -9948,6 +10018,13 @@ async def test_doubly_nested_graph_state(
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": AnyStr(),
|
||||
"checkpoint_id": AnyStr(),
|
||||
"checkpoint_map": AnyDict(
|
||||
{
|
||||
"": AnyStr(),
|
||||
AnyStr("child:"): AnyStr(),
|
||||
AnyStr(re.compile(r"child:.+|child1:")): AnyStr(),
|
||||
}
|
||||
),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
@@ -9967,6 +10044,19 @@ async def test_doubly_nested_graph_state(
|
||||
),
|
||||
]
|
||||
|
||||
# replay grandchild checkpoint
|
||||
assert [
|
||||
c async for c in app.astream(None, grandchild_history[2].config, subgraphs=True)
|
||||
] == [
|
||||
(
|
||||
(AnyStr("child:"), AnyStr("child_1:")),
|
||||
{
|
||||
"grandchild_1": {"my_key": "hi my value here"},
|
||||
"__metadata__": {"cached": True},
|
||||
},
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.skip("TODO")
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
|
||||
|
||||
Reference in New Issue
Block a user