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

Merge pull request #1538 from langchain-ai/nc/29aug/nested-replay

Browse Source 
Replay a nested checkpoint
This commit is contained in:
Nuno Campos
2024-08-30 10:30:14 -07:00
committed by GitHub
parent 30598cd913 4fc4080597
commit 24a7e6403f
9 changed files with 495 additions and 3788 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/__init__.py
+73 -42
View File
@@ -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", {}):
libs/langgraph/langgraph/pregel/algo.py
+13 -3
View File
@@ -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,
)
)
libs/langgraph/langgraph/pregel/config.py
+12
View File
@@ -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}}
libs/langgraph/langgraph/pregel/debug.py
+16 -12
View File
@@ -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],
libs/langgraph/langgraph/pregel/loop.py
+34 -5
View File
@@ -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, [])
libs/langgraph/langgraph/pregel/types.py
+10
View File
@@ -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
libs/langgraph/tests/any_str.py
+9 -4
View File
@@ -1,16 +1,21 @@
from typing import Any, Sequence
import re
from typing import Any, Sequence, Union
class AnyStr(str):
def __init__(self, prefix: str = "") -> None:
def __init__(self, prefix: Union[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):
libs/langgraph/tests/test_pregel.py
+87 -1816
View File
File diff suppressed because it is too large Load Diff
libs/langgraph/tests/test_pregel_async.py
+241 -1906
View File
File diff suppressed because it is too large Load Diff