WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-07 10:17:50 +02:00
Code Issues Packages Projects Releases Wiki Activity

Implement get_state and get_state_history for nested graphs

Browse Source
This commit is contained in:
Nuno Campos
2024-08-29 09:45:45 -07:00
parent c4010c03a6
commit ca8614f873
16 changed files with 2833 additions and 1603 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint/langgraph/checkpoint/base/__init__.py
+3 -3
View File
@@ -51,10 +51,10 @@ class CheckpointMetadata(TypedDict, total=False):
Mapping from node name to writes emitted by that node.
"""
score: Optional[int]
"""The score of the checkpoint.
parents: dict[str, str]
"""The IDs of the parent checkpoints.
The score can be used to mark a checkpoint as "good".
Mapping from checkpoint namespace to checkpoint ID.
"""
libs/langgraph/langgraph/constants.py
+4 -2
View File
@@ -5,6 +5,7 @@ INPUT = "__input__"
CONFIG_KEY_SEND = "__pregel_send"
CONFIG_KEY_READ = "__pregel_read"
CONFIG_KEY_CHECKPOINTER = "__pregel_checkpointer"
CONFIG_KEY_CHECKPOINT_MAP = "checkpoint_map"
CONFIG_KEY_STORE = "__pregel_store"
CONFIG_KEY_RESUMING = "__pregel_resuming"
CONFIG_KEY_TASK_ID = "__pregel_task_id"
@@ -19,6 +20,7 @@ RESERVED = {
CONFIG_KEY_SEND,
CONFIG_KEY_READ,
CONFIG_KEY_CHECKPOINTER,
CONFIG_KEY_CHECKPOINT_MAP,
CONFIG_KEY_STORE,
CONFIG_KEY_RESUMING,
CONFIG_KEY_TASK_ID,
@@ -30,8 +32,8 @@ TAG_HIDDEN = "langsmith:hidden"
START = "__start__"
END = "__end__"
CHECKPOINT_NAMESPACE_SEPARATOR = "|"
SEND_CHECKPOINT_NAMESPACE_SEPARATOR = ":"
NS_SEP = "|"
NS_END = ":"
class Send:
libs/langgraph/langgraph/graph/graph.py
+3 -6
View File
@@ -26,9 +26,9 @@ from langchain_core.runnables.graph import Node as DrawableNode
from langgraph.channels.ephemeral_value import EphemeralValue
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.constants import (
CHECKPOINT_NAMESPACE_SEPARATOR,
END,
SEND_CHECKPOINT_NAMESPACE_SEPARATOR,
NS_END,
NS_SEP,
START,
TAG_HIDDEN,
Send,
@@ -160,10 +160,7 @@ class Graph:
metadata: Optional[dict[str, Any]] = None,
) -> None:
if isinstance(node, str):
for character in (
CHECKPOINT_NAMESPACE_SEPARATOR,
SEND_CHECKPOINT_NAMESPACE_SEPARATOR,
):
for character in (NS_SEP, NS_END):
if character in node:
raise ValueError(
f"'{character}' is a reserved character and is not allowed in the node names."
libs/langgraph/langgraph/graph/state.py
+2 -9
View File
@@ -31,11 +31,7 @@ from langgraph.channels.ephemeral_value import EphemeralValue
from langgraph.channels.last_value import LastValue
from langgraph.channels.named_barrier_value import NamedBarrierValue
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.constants import (
CHECKPOINT_NAMESPACE_SEPARATOR,
SEND_CHECKPOINT_NAMESPACE_SEPARATOR,
TAG_HIDDEN,
)
from langgraph.constants import NS_END, NS_SEP, TAG_HIDDEN
from langgraph.errors import InvalidUpdateError
from langgraph.graph.graph import (
END,
@@ -321,10 +317,7 @@ class StateGraph(Graph):
if node == END or node == START:
raise ValueError(f"Node `{node}` is reserved.")
for character in (
CHECKPOINT_NAMESPACE_SEPARATOR,
SEND_CHECKPOINT_NAMESPACE_SEPARATOR,
):
for character in (NS_SEP, NS_END):
if character in node:
raise ValueError(
f"'{character}' is a reserved character and is not allowed in the node names."
libs/langgraph/langgraph/pregel/__init__.py
+275 -281
View File
@@ -5,7 +5,6 @@ import concurrent.futures
import time
from collections import deque
from functools import partial
from inspect import signature
from typing import (
Any,
AsyncIterator,
@@ -60,14 +59,14 @@ 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,
ERROR,
INTERRUPT,
SEND_CHECKPOINT_NAMESPACE_SEPARATOR,
NS_END,
NS_SEP,
)
from langgraph.errors import GraphInterrupt, GraphRecursionError, InvalidUpdateError
from langgraph.managed.base import ManagedValueSpec
@@ -83,7 +82,6 @@ from langgraph.pregel.debug import (
print_step_writes,
tasks_w_writes,
)
from langgraph.pregel.get_state import assemble_state_snapshot_hierarchy
from langgraph.pregel.io import read_channels
from langgraph.pregel.loop import AsyncPregelLoop, SyncPregelLoop
from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
@@ -182,181 +180,6 @@ class Channel:
)
def _get_checkpoint_ns_to_graph(
graph: Pregel,
checkpoint_ns_to_graph: Optional[dict[str, Pregel]] = None,
checkpoint_ns: str = "",
max_depth: int = 10,
) -> Pregel:
if checkpoint_ns_to_graph is None:
checkpoint_ns_to_graph = {}
if max_depth <= 0:
raise RecursionError(
"Reached maximum recursion depth while building checkpoint NS -> graph mapping."
)
for node_name, node in graph.nodes.items():
new_checkpoint_ns = (
f"{checkpoint_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{node_name}"
if checkpoint_ns
else node_name
)
if isinstance(node.bound, Pregel):
_get_checkpoint_ns_to_graph(
node.bound, checkpoint_ns_to_graph, new_checkpoint_ns, max_depth - 1
)
elif isinstance(node.bound, RunnableSequence):
for runnable in node.bound.steps:
if isinstance(runnable, Pregel):
_get_checkpoint_ns_to_graph(
runnable,
checkpoint_ns_to_graph,
new_checkpoint_ns,
max_depth - 1,
)
checkpoint_ns_to_graph[checkpoint_ns] = graph
return checkpoint_ns_to_graph
def _prepare_state_snapshot(
config: RunnableConfig,
checkpoint_ns_to_graph: dict[str, Pregel],
checkpoint_tuples: Iterator[CheckpointTuple],
) -> StateSnapshot:
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = config["configurable"].get("checkpoint_id")
checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot] = {}
for saved in checkpoint_tuples:
saved_checkpoint_ns = saved.config["configurable"]["checkpoint_ns"]
saved_checkpoint_id = saved.config["configurable"]["checkpoint_id"]
if checkpoint_id and saved_checkpoint_id != checkpoint_id:
continue
graph_checkpoint_ns = saved_checkpoint_ns.split(
SEND_CHECKPOINT_NAMESPACE_SEPARATOR
)[0]
graph = checkpoint_ns_to_graph.get(graph_checkpoint_ns)
if graph is None:
continue
with ChannelsManager(
graph.channels, saved.checkpoint, saved.config, skip_context=True
) as (
channels,
managed,
):
next_tasks = prepare_next_tasks(
saved.checkpoint,
graph.nodes,
channels,
managed,
saved.config,
saved.metadata.get("step", -1) + 1,
for_execution=False,
)
state_snapshot = StateSnapshot(
read_channels(channels, graph.stream_channels_asis),
tuple(t.name for t in next_tasks),
saved.config,
saved.metadata,
saved.checkpoint["ts"],
saved.parent_config,
tasks_w_writes(next_tasks, saved.pending_writes),
)
checkpoint_ns_to_state_snapshots[saved_checkpoint_ns] = state_snapshot
if not checkpoint_ns_to_state_snapshots:
return StateSnapshot(
values={},
next=(),
config=config,
metadata=None,
created_at=None,
parent_config=None,
tasks=(),
)
state_snapshot = assemble_state_snapshot_hierarchy(
checkpoint_ns, checkpoint_ns_to_state_snapshots
)
return state_snapshot
async def _prepare_state_snapshot_async(
config: RunnableConfig,
checkpoint_ns_to_graph: dict[str, Pregel],
checkpoint_tuples: AsyncIterator[CheckpointTuple],
) -> StateSnapshot:
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_id = config["configurable"].get("checkpoint_id")
checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot] = {}
async for saved in checkpoint_tuples:
saved_checkpoint_ns = saved.config["configurable"]["checkpoint_ns"]
saved_checkpoint_id = saved.config["configurable"]["checkpoint_id"]
if checkpoint_id and saved_checkpoint_id != checkpoint_id:
continue
graph_checkpoint_ns = saved_checkpoint_ns.split(
SEND_CHECKPOINT_NAMESPACE_SEPARATOR
)[0]
graph = checkpoint_ns_to_graph.get(graph_checkpoint_ns)
if graph is None:
continue
async with AsyncChannelsManager(
graph.channels, saved.checkpoint, saved.config, skip_context=True
) as (channels, managed):
next_tasks = prepare_next_tasks(
saved.checkpoint,
graph.nodes,
channels,
managed,
saved.config,
saved.metadata.get("step", -1) + 1,
for_execution=False,
)
state_snapshot = StateSnapshot(
read_channels(channels, graph.stream_channels_asis),
tuple(t.name for t in next_tasks),
saved.config,
saved.metadata,
saved.checkpoint["ts"],
saved.parent_config,
tasks_w_writes(next_tasks, saved.pending_writes),
)
checkpoint_ns_to_state_snapshots[saved_checkpoint_ns] = state_snapshot
if not checkpoint_ns_to_state_snapshots:
return StateSnapshot(
values={},
next=(),
config=config,
metadata=None,
created_at=None,
parent_config=None,
tasks=(),
)
state_snapshot = assemble_state_snapshot_hierarchy(
checkpoint_ns, checkpoint_ns_to_state_snapshots
)
return state_snapshot
def _has_nested_interrupts(
graph: Pregel,
) -> bool:
for child in graph.subgraphs:
if child.interrupt_after_nodes or child.interrupt_before_nodes:
return True
else:
return False
class Pregel(
RunnableSerializable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]
):
@@ -526,60 +349,211 @@ class Pregel(
k for k in self.channels if isinstance(self.channels[k], BaseChannel)
]
@property
def subgraphs(self) -> Iterator[Pregel]:
for node in self.nodes.values():
def get_subgraphs(self, recurse: bool = False) -> Iterator[tuple[str, Pregel]]:
for name, node in self.nodes.items():
# find the subgraph, if any
graph: Optional[Pregel] = None
if isinstance(node.bound, Pregel):
yield node.bound
yield from node.bound.subgraphs
graph = node.bound
elif isinstance(node.bound, RunnableSequence):
for runnable in node.bound.steps:
if isinstance(runnable, Pregel):
yield runnable
yield from runnable.subgraphs
graph = runnable
break
# if found, yield recursively
if graph:
yield name, graph
if recurse:
yield from (
(f"{name}{NS_SEP}{n}", s)
for n, s in graph.get_subgraphs(recurse=recurse)
)
def get_state(self, config: RunnableConfig) -> StateSnapshot:
async def aget_subgraphs(
self, recursive: bool = False
) -> AsyncIterator[tuple[str, Pregel]]:
for name, node in self.get_subgraphs(recurse=recursive):
yield name, node
def _prepare_state_snapshot(
self,
config: RunnableConfig,
saved: Optional[CheckpointTuple],
recurse: Optional[BaseCheckpointSaver] = False,
) -> StateSnapshot:
if not saved:
return StateSnapshot(
values={},
next=(),
config=config,
metadata=None,
created_at=None,
parent_config=None,
tasks=(),
)
with ChannelsManager(
self.channels, saved.checkpoint, saved.config, skip_context=True
) as (channels, managed):
# tasks for this checkpoint
next_tasks = prepare_next_tasks(
saved.checkpoint,
self.nodes,
channels,
managed,
saved.config,
saved.metadata.get("step", -1) + 1,
for_execution=False,
)
# get the subgraphs
subgraphs = dict(self.get_subgraphs())
parent_ns = saved.config["configurable"].get("checkpoint_ns", "")
task_states: dict[str, Union[RunnableConfig, StateSnapshot]] = {}
for task in next_tasks:
if task.name not in subgraphs:
continue
# assemble checkpoint_ns for this task
task_ns = f"{task.name}{NS_END}{task.id}"
if parent_ns:
task_ns = f"{parent_ns}{NS_SEP}{task_ns}"
if not recurse:
# set config as signal that subgraph checkpoints exist
config = {
"configurable": {
"thread_id": saved.config["configurable"]["thread_id"],
"checkpoint_ns": task_ns,
}
}
task_states[task.id] = config
else:
# get the state of the subgraph
config = {
"configurable": {
CONFIG_KEY_CHECKPOINTER: recurse,
"thread_id": saved.config["configurable"]["thread_id"],
"checkpoint_ns": task_ns,
}
}
task_states[task.id] = subgraphs[task.name].get_state(
config, subgraphs=True
)
# assemble the state snapshot
return StateSnapshot(
read_channels(channels, self.stream_channels_asis),
tuple(t.name for t in next_tasks),
saved.config,
saved.metadata,
saved.checkpoint["ts"],
saved.parent_config,
tasks_w_writes(next_tasks, saved.pending_writes, task_states),
)
async def _aprepare_state_snapshot(
self,
config: RunnableConfig,
saved: Optional[CheckpointTuple],
recurse: Optional[BaseCheckpointSaver] = False,
) -> StateSnapshot:
if not saved:
return StateSnapshot(
values={},
next=(),
config=config,
metadata=None,
created_at=None,
parent_config=None,
tasks=(),
)
async with AsyncChannelsManager(
self.channels, saved.checkpoint, saved.config, skip_context=True
) as (
channels,
managed,
):
# tasks for this checkpoint
next_tasks = prepare_next_tasks(
saved.checkpoint,
self.nodes,
channels,
managed,
saved.config,
saved.metadata.get("step", -1) + 1,
for_execution=False,
)
# get the subgraphs
subgraphs = dict(self.get_subgraphs())
parent_ns = saved.config["configurable"].get("checkpoint_ns", "")
task_states: dict[str, Union[RunnableConfig, StateSnapshot]] = {}
for task in next_tasks:
if task.name not in subgraphs:
continue
# assemble checkpoint_ns for this task
task_ns = f"{task.name}{NS_END}{task.id}"
if parent_ns:
task_ns = f"{parent_ns}{NS_SEP}{task_ns}"
if not recurse:
# set config as signal that subgraph checkpoints exist
config = {
"configurable": {
"thread_id": saved.config["configurable"]["thread_id"],
"checkpoint_ns": task_ns,
}
}
task_states[task.id] = config
else:
# get the state of the subgraph
config = {
"configurable": {
CONFIG_KEY_CHECKPOINTER: recurse,
"thread_id": saved.config["configurable"]["thread_id"],
"checkpoint_ns": task_ns,
}
}
task_states[task.id] = await subgraphs[task.name].aget_state(
config, subgraphs=recurse
)
# assemble the state snapshot
return StateSnapshot(
read_channels(channels, self.stream_channels_asis),
tuple(t.name for t in next_tasks),
saved.config,
saved.metadata,
saved.checkpoint["ts"],
saved.parent_config,
tasks_w_writes(next_tasks, saved.pending_writes, task_states),
)
def get_state(
self, config: RunnableConfig, *, subgraphs: bool = False
) -> StateSnapshot:
"""Get the current state of the graph."""
if not self.checkpointer:
checkpointer: Optional[BaseCheckpointSaver] = config["configurable"].get(
CONFIG_KEY_CHECKPOINTER, self.checkpointer
)
if not checkpointer:
raise ValueError("No checkpointer set")
config = merge_configs(self.config, config) if self.config else config
saved = self.checkpointer.get_tuple(config)
checkpoint_config = saved.config if saved else config
checkpoint_ns_to_graph: dict[str, Pregel] = _get_checkpoint_ns_to_graph(self)
# we only lookup subgraph checkpoints if we actually have subgraphs
if len(set(checkpoint_ns_to_graph)) == 1:
checkpoint_tuples = (saved,)
else:
checkpoint_tuples = self.checkpointer.list(saved.config)
return _prepare_state_snapshot(
checkpoint_config, checkpoint_ns_to_graph, checkpoint_tuples
saved = checkpointer.get_tuple(config)
return self._prepare_state_snapshot(
config, saved, recurse=checkpointer if subgraphs else None
)
async def aget_state(self, config: RunnableConfig) -> StateSnapshot:
async def aget_state(
self, config: RunnableConfig, *, subgraphs: bool = False
) -> StateSnapshot:
"""Get the current state of the graph."""
if not self.checkpointer:
checkpointer: Optional[BaseCheckpointSaver] = config["configurable"].get(
CONFIG_KEY_CHECKPOINTER, self.checkpointer
)
if not checkpointer:
raise ValueError("No checkpointer set")
config = merge_configs(self.config, config) if self.config else config
saved = await self.checkpointer.aget_tuple(config)
checkpoint_config = saved.config if saved else config
checkpoint_ns_to_graph: dict[str, Pregel] = _get_checkpoint_ns_to_graph(self)
# we only lookup subgraph checkpoints if we actually have subgraphs
if len(set(checkpoint_ns_to_graph)) == 1:
async def alist_checkpoints():
yield saved
checkpoint_tuples = alist_checkpoints()
else:
checkpoint_tuples = self.checkpointer.alist(saved.config)
return await _prepare_state_snapshot_async(
checkpoint_config, checkpoint_ns_to_graph, checkpoint_tuples
saved = await checkpointer.aget_tuple(config)
return await self._aprepare_state_snapshot(
config, saved, recurse=checkpointer if subgraphs else None
)
def get_state_history(
@@ -591,37 +565,47 @@ class Pregel(
limit: Optional[int] = None,
) -> Iterator[StateSnapshot]:
"""Get the history of the state of the graph."""
if not self.checkpointer:
checkpointer: Optional[BaseCheckpointSaver] = config["configurable"].get(
CONFIG_KEY_CHECKPOINTER, self.checkpointer
)
if not checkpointer:
raise ValueError("No checkpointer set")
if (
filter is not None
and signature(self.checkpointer.list).parameters.get("filter") is None
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:
yield from pregel.get_state_history(
{
"configurable": {
**config["configurable"],
CONFIG_KEY_CHECKPOINTER: checkpointer,
}
},
filter=filter,
before=before,
limit=limit,
)
return
else:
raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
config = merge_configs(
self.config, config, {"configurable": {"checkpoint_ns": checkpoint_ns}}
)
# eagerly consume list() to avoid holding up the db cursor
for checkpoint_tuple in list(
checkpointer.list(config, before=before, limit=limit, filter=filter)
):
raise ValueError("Checkpointer does not support filtering")
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_ns_to_graph = _get_checkpoint_ns_to_graph(self)
# find all matching checkpoint tuples for parent and subgraphs
checkpoint_tuples = [
checkpoint_tuple
for checkpoint_tuple in self.checkpointer.list(
merge_configs(self.config, config) if self.config else config,
before=before,
limit=limit,
filter=filter,
yield self._prepare_state_snapshot(
checkpoint_tuple.config, checkpoint_tuple
)
]
for checkpoint_tuple in checkpoint_tuples:
if (
checkpoint_tuple.config["configurable"]["checkpoint_ns"]
!= checkpoint_ns
):
continue
state_snapshot = _prepare_state_snapshot(
checkpoint_tuple.config, checkpoint_ns_to_graph, iter(checkpoint_tuples)
)
yield state_snapshot
async def aget_state_history(
self,
@@ -632,43 +616,51 @@ class Pregel(
limit: Optional[int] = None,
) -> AsyncIterator[StateSnapshot]:
"""Get the history of the state of the graph."""
if not self.checkpointer:
checkpointer: Optional[BaseCheckpointSaver] = config["configurable"].get(
CONFIG_KEY_CHECKPOINTER, self.checkpointer
)
if not checkpointer:
raise ValueError("No checkpointer set")
if (
filter is not None
and signature(self.checkpointer.list).parameters.get("filter") is None
):
raise ValueError("Checkpointer does not support filtering")
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
checkpoint_ns_to_graph = _get_checkpoint_ns_to_graph(self)
# find all matching checkpoint tuples for parent and subgraphs
checkpoint_tuples = [
checkpoint_tuple
async for checkpoint_tuple in self.checkpointer.alist(
merge_configs(self.config, config) if self.config else config,
before=before,
limit=limit,
filter=filter,
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:
async for state in pregel.aget_state_history(
{
"configurable": {
**config["configurable"],
CONFIG_KEY_CHECKPOINTER: checkpointer,
}
},
filter=filter,
before=before,
limit=limit,
):
yield state
return
else:
raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
# turn matching checkpoint tuples into an async iterator
async def alist_checkpoints() -> AsyncIterator[CheckpointTuple]:
for checkpoint_tuple in checkpoint_tuples:
yield checkpoint_tuple
for checkpoint_tuple in checkpoint_tuples:
if (
checkpoint_tuple.config["configurable"]["checkpoint_ns"]
!= checkpoint_ns
):
continue
state_snapshot = await _prepare_state_snapshot_async(
checkpoint_tuple.config, checkpoint_ns_to_graph, alist_checkpoints()
config = merge_configs(
self.config, config, {"configurable": {"checkpoint_ns": checkpoint_ns}}
)
# eagerly consume list() to avoid holding up the db cursor
for checkpoint_tuple in [
c
async for c in checkpointer.alist(
config, before=before, limit=limit, filter=filter
)
]:
yield await self._aprepare_state_snapshot(
checkpoint_tuple.config, checkpoint_tuple
)
yield state_snapshot
def update_state(
self,
@@ -716,6 +708,7 @@ class Pregel(
"source": "update",
"step": step + 1,
"writes": {},
"parents": saved.metadata.get("parents", {}) if saved else {},
},
{},
)
@@ -805,6 +798,7 @@ class Pregel(
"source": "update",
"step": step + 1,
"writes": {as_node: values},
"parents": saved.metadata.get("parents", {}) if saved else {},
},
get_new_channel_versions(
checkpoint_previous_versions, checkpoint["channel_versions"]
@@ -853,6 +847,7 @@ class Pregel(
"source": "update",
"step": step + 1,
"writes": {},
"parents": saved.metadata.get("parents", {}) if saved else {},
},
{},
)
@@ -942,6 +937,7 @@ class Pregel(
"source": "update",
"step": step + 1,
"writes": {as_node: values},
"parents": saved.metadata.get("parents", {}) if saved else {},
},
get_new_channel_versions(
checkpoint_previous_versions, checkpoint["channel_versions"]
@@ -979,10 +975,8 @@ class Pregel(
if config and config.get("configurable", {}).get(CONFIG_KEY_READ) is not None:
# if being called as a node in another graph, always use values mode
stream_mode = ["values"]
if (
config is not None
and config.get("configurable", {}).get(CONFIG_KEY_CHECKPOINTER)
and (interrupt_after or interrupt_before or _has_nested_interrupts(self))
if config is not None and config.get("configurable", {}).get(
CONFIG_KEY_CHECKPOINTER
):
checkpointer: Optional[BaseCheckpointSaver] = config["configurable"][
CONFIG_KEY_CHECKPOINTER
libs/langgraph/langgraph/pregel/algo.py
+23 -17
View File
@@ -31,13 +31,14 @@ from langgraph.checkpoint.base import (
create_checkpoint,
)
from langgraph.constants import (
CHECKPOINT_NAMESPACE_SEPARATOR,
CONFIG_KEY_CHECKPOINT_MAP,
CONFIG_KEY_CHECKPOINTER,
CONFIG_KEY_READ,
CONFIG_KEY_RESUMING,
CONFIG_KEY_SEND,
CONFIG_KEY_TASK_ID,
INTERRUPT,
NS_SEP,
RESERVED,
TAG_HIDDEN,
TASKS,
@@ -272,7 +273,8 @@ def prepare_next_tasks(
checkpointer: Optional[BaseCheckpointSaver] = None,
manager: Union[None, ParentRunManager, AsyncParentRunManager] = None,
) -> Union[list[PregelTask], list[PregelExecutableTask]]:
parent_ns = config.get("configurable", {}).get("checkpoint_ns", "")
configurable = config.get("configurable", {})
parent_ns = configurable.get("checkpoint_ns", "")
tasks: Union[list[PregelTask], list[PregelExecutableTask]] = []
# Consume pending packets
for packet in checkpoint["pending_sends"]:
@@ -291,9 +293,7 @@ def prepare_next_tasks(
"langgraph_task_idx": len(tasks),
}
checkpoint_ns = (
f"{parent_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{packet.node}"
if parent_ns
else packet.node
f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node
)
task_id = str(
uuid5(UUID(checkpoint["id"]), json.dumps((checkpoint_ns, metadata)))
@@ -341,9 +341,16 @@ def prepare_next_tasks(
PregelTaskWrites(packet.node, writes, triggers),
config,
),
CONFIG_KEY_CHECKPOINTER: checkpointer,
CONFIG_KEY_CHECKPOINTER: (
checkpointer
or configurable.get(CONFIG_KEY_CHECKPOINTER)
),
CONFIG_KEY_CHECKPOINT_MAP: {
**configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {}),
parent_ns: checkpoint["id"],
},
CONFIG_KEY_RESUMING: is_resuming,
"checkpoint_id": checkpoint["id"],
"checkpoint_id": None,
"checkpoint_ns": f"{checkpoint_ns}:{task_id}",
},
),
@@ -388,11 +395,7 @@ def prepare_next_tasks(
"langgraph_triggers": triggers,
"langgraph_task_idx": len(tasks),
}
checkpoint_ns = (
f"{parent_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{name}"
if parent_ns
else name
)
checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name
task_id = str(
uuid5(
UUID(checkpoint["id"]),
@@ -443,13 +446,16 @@ def prepare_next_tasks(
),
CONFIG_KEY_CHECKPOINTER: (
checkpointer
or config["configurable"].get(
CONFIG_KEY_CHECKPOINTER
)
or configurable.get(CONFIG_KEY_CHECKPOINTER)
),
CONFIG_KEY_CHECKPOINT_MAP: {
**configurable.get(
CONFIG_KEY_CHECKPOINT_MAP, {}
),
parent_ns: checkpoint["id"],
},
CONFIG_KEY_RESUMING: is_resuming,
"checkpoint_id": checkpoint["id"],
"checkpoint_ns": checkpoint_ns,
"checkpoint_ns": f"{checkpoint_ns}:{task_id}",
},
),
triggers,
libs/langgraph/langgraph/pregel/debug.py
+4 -2
View File
@@ -13,7 +13,7 @@ from langgraph.channels.base import BaseChannel
from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, PendingWrite
from langgraph.constants import ERROR, INTERRUPT, TAG_HIDDEN
from langgraph.pregel.io import read_channels
from langgraph.pregel.types import PregelExecutableTask, PregelTask
from langgraph.pregel.types import PregelExecutableTask, PregelTask, StateSnapshot
class TaskPayload(TypedDict):
@@ -160,7 +160,7 @@ def map_debug_checkpoint(
"name": t.name,
"interrupts": tuple(asdict(i) for i in t.interrupts),
}
for t in tasks_w_writes(tasks, pending_writes)
for t in tasks_w_writes(tasks, pending_writes, None)
],
},
}
@@ -215,6 +215,7 @@ def print_step_checkpoint(
def tasks_w_writes(
tasks: list[PregelExecutableTask],
pending_writes: Optional[list[PendingWrite]],
states: Optional[dict[str, Union[RunnableConfig, StateSnapshot]]],
) -> tuple[PregelTask, ...]:
pending_writes = pending_writes or []
return tuple(
@@ -232,6 +233,7 @@ def tasks_w_writes(
tuple(
v for tid, n, v in pending_writes if tid == task.id and n == INTERRUPT
),
states.get(task.id) if states else None,
)
for task in tasks
)
libs/langgraph/langgraph/pregel/get_state.py
+6 -6
View File
@@ -1,4 +1,4 @@
from langgraph.constants import CHECKPOINT_NAMESPACE_SEPARATOR
from langgraph.constants import NS_SEP
from langgraph.pregel.types import StateSnapshot
@@ -8,26 +8,26 @@ def assemble_state_snapshot_hierarchy(
) -> StateSnapshot:
checkpoint_ns_list_to_visit = sorted(
checkpoint_ns_to_state_snapshots.keys(),
key=lambda x: len(x.split(CHECKPOINT_NAMESPACE_SEPARATOR)),
key=lambda x: len(x.split(NS_SEP)),
)
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)
*path, subgraph_node = checkpoint_ns.split(NS_SEP)
parent_checkpoint_ns = NS_SEP.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 {}),
**(parent_state_snapshot.subgraphs or {}),
subgraph_node: state_snapshot,
}
checkpoint_ns_to_state_snapshots[parent_checkpoint_ns] = (
checkpoint_ns_to_state_snapshots[
parent_checkpoint_ns
]._replace(subgraph_state_snapshots=parent_subgraph_snapshots)
]._replace(subgraphs=parent_subgraph_snapshots)
)
state_snapshot = checkpoint_ns_to_state_snapshots.pop(root_checkpoint_ns, None)
libs/langgraph/langgraph/pregel/loop.py
+5 -10
View File
@@ -36,6 +36,7 @@ from langgraph.checkpoint.base import (
empty_checkpoint,
)
from langgraph.constants import (
CONFIG_KEY_CHECKPOINT_MAP,
CONFIG_KEY_READ,
CONFIG_KEY_RESUMING,
ERROR,
@@ -361,21 +362,15 @@ class PregelLoop:
def _put_checkpoint(self, metadata: CheckpointMetadata) -> None:
# assign step
metadata["step"] = self.step
metadata["parents"] = self.config["configurable"].get(
CONFIG_KEY_CHECKPOINT_MAP, {}
)
# bail if no checkpointer
if self._checkpointer_put_after_previous is not None:
# create new checkpoint
self.checkpoint_metadata = metadata
self.checkpoint = create_checkpoint(
self.checkpoint,
self.channels,
self.step,
# child graphs keep at most one checkpoint per parent checkpoint
# this is achieved by writing child checkpoints as progress is made
# (so that error recovery / resuming from interrupt don't lose work)
# but doing so always with an id equal to that of the parent checkpoint
id=self.config["configurable"]["checkpoint_id"]
if self.is_nested
else None,
self.checkpoint, self.channels, self.step
)
self.checkpoint_config = {
libs/langgraph/langgraph/pregel/metadata.py
View File
libs/langgraph/langgraph/pregel/types.py
+1 -2
View File
@@ -62,6 +62,7 @@ class PregelTask(NamedTuple):
name: str
error: Optional[Exception] = None
interrupts: tuple[Interrupt, ...] = ()
state: Union[None, RunnableConfig, "StateSnapshot"] = None
class PregelExecutableTask(NamedTuple):
@@ -92,8 +93,6 @@ class StateSnapshot(NamedTuple):
"""Config used to fetch the parent snapshot, if any"""
tasks: tuple[PregelTask, ...]
"""Tasks to execute in this step. If already attempted, may contain an error."""
subgraph_state_snapshots: Optional[dict[str, "StateSnapshot"]] = None
"""State snapshots of subgraphs represented as a mapping from checkpoint namespace (`checkpoint_ns`) to snapshot."""
All = Literal["*"]
libs/langgraph/tests/__snapshots__/test_pregel.ambr
+110 -80
View File
@@ -805,6 +805,7 @@
"name": "tools"
},
"metadata": {
"parents": {},
"version": 2,
"variant": "b"
}
@@ -856,7 +857,8 @@
graph TD;
__start__([__start__]):::first
agent(agent)
tools(tools<hr/><small><em>version = 2
tools(tools<hr/><small><em>parents = {}
version = 2
variant = b</em></small>)
__end__([__end__]):::last
__start__ --> agent;
@@ -952,6 +954,7 @@
"name": "tools"
},
"metadata": {
"parents": {},
"version": 2,
"variant": "b"
}
@@ -1084,6 +1087,8 @@
}),
'id': 'tools',
'metadata': dict({
'parents': dict({
}),
'variant': 'b',
'version': 2,
}),
@@ -1103,7 +1108,8 @@
graph TD;
__start__([__start__]):::first
agent(agent<hr/><small><em>__interrupt = after</em></small>)
tools(tools<hr/><small><em>version = 2
tools(tools<hr/><small><em>parents = {}
version = 2
variant = b</em></small>)
__end__([__end__]):::last
__start__ --> agent;
@@ -1150,6 +1156,7 @@
"name": "tools"
},
"metadata": {
"parents": {},
"version": 2,
"variant": "b"
}
@@ -1201,7 +1208,8 @@
graph TD;
__start__([__start__]):::first
agent(agent)
tools(tools<hr/><small><em>version = 2
tools(tools<hr/><small><em>parents = {}
version = 2
variant = b</em></small>)
__end__([__end__]):::last
__start__ --> agent;
@@ -1297,6 +1305,7 @@
"name": "tools"
},
"metadata": {
"parents": {},
"version": 2,
"variant": "b"
}
@@ -1429,6 +1438,8 @@
}),
'id': 'tools',
'metadata': dict({
'parents': dict({
}),
'variant': 'b',
'version': 2,
}),
@@ -1448,7 +1459,8 @@
graph TD;
__start__([__start__]):::first
agent(agent<hr/><small><em>__interrupt = after</em></small>)
tools(tools<hr/><small><em>version = 2
tools(tools<hr/><small><em>parents = {}
version = 2
variant = b</em></small>)
__end__([__end__]):::last
__start__ --> agent;
@@ -1495,6 +1507,7 @@
"name": "tools"
},
"metadata": {
"parents": {},
"version": 2,
"variant": "b"
}
@@ -1546,7 +1559,8 @@
graph TD;
__start__([__start__]):::first
agent(agent)
tools(tools<hr/><small><em>version = 2
tools(tools<hr/><small><em>parents = {}
version = 2
variant = b</em></small>)
__end__([__end__]):::last
__start__ --> agent;
@@ -1642,6 +1656,7 @@
"name": "tools"
},
"metadata": {
"parents": {},
"version": 2,
"variant": "b"
}
@@ -1774,6 +1789,8 @@
}),
'id': 'tools',
'metadata': dict({
'parents': dict({
}),
'variant': 'b',
'version': 2,
}),
@@ -1793,7 +1810,8 @@
graph TD;
__start__([__start__]):::first
agent(agent<hr/><small><em>__interrupt = after</em></small>)
tools(tools<hr/><small><em>version = 2
tools(tools<hr/><small><em>parents = {}
version = 2
variant = b</em></small>)
__end__([__end__]):::last
__start__ --> agent;
@@ -1840,6 +1858,7 @@
"name": "tools"
},
"metadata": {
"parents": {},
"version": 2,
"variant": "b"
}
@@ -1891,7 +1910,8 @@
graph TD;
__start__([__start__]):::first
agent(agent)
tools(tools<hr/><small><em>version = 2
tools(tools<hr/><small><em>parents = {}
version = 2
variant = b</em></small>)
__end__([__end__]):::last
__start__ --> agent;
@@ -1987,6 +2007,7 @@
"name": "tools"
},
"metadata": {
"parents": {},
"version": 2,
"variant": "b"
}
@@ -2119,6 +2140,8 @@
}),
'id': 'tools',
'metadata': dict({
'parents': dict({
}),
'variant': 'b',
'version': 2,
}),
@@ -2138,7 +2161,8 @@
graph TD;
__start__([__start__]):::first
agent(agent<hr/><small><em>__interrupt = after</em></small>)
tools(tools<hr/><small><em>version = 2
tools(tools<hr/><small><em>parents = {}
version = 2
variant = b</em></small>)
__end__([__end__]):::last
__start__ --> agent;
@@ -2185,6 +2209,7 @@
"name": "tools"
},
"metadata": {
"parents": {},
"version": 2,
"variant": "b"
}
@@ -2236,7 +2261,8 @@
graph TD;
__start__([__start__]):::first
agent(agent)
tools(tools<hr/><small><em>version = 2
tools(tools<hr/><small><em>parents = {}
version = 2
variant = b</em></small>)
__end__([__end__]):::last
__start__ --> agent;
@@ -2332,6 +2358,7 @@
"name": "tools"
},
"metadata": {
"parents": {},
"version": 2,
"variant": "b"
}
@@ -2464,6 +2491,8 @@
}),
'id': 'tools',
'metadata': dict({
'parents': dict({
}),
'variant': 'b',
'version': 2,
}),
@@ -2483,7 +2512,8 @@
graph TD;
__start__([__start__]):::first
agent(agent<hr/><small><em>__interrupt = after</em></small>)
tools(tools<hr/><small><em>version = 2
tools(tools<hr/><small><em>parents = {}
version = 2
variant = b</em></small>)
__end__([__end__]):::last
__start__ --> agent;
@@ -2658,76 +2688,6 @@
'''
# ---
# name: test_conditional_state_graph_with_list_edge_inputs
'''
{
"nodes": [
{
"id": "__start__",
"type": "schema",
"data": "__start__"
},
{
"id": "A",
"type": "runnable",
"data": {
"id": [
"langgraph",
"utils",
"RunnableCallable"
],
"name": "A"
}
},
{
"id": "B",
"type": "runnable",
"data": {
"id": [
"langgraph",
"utils",
"RunnableCallable"
],
"name": "B"
}
},
{
"id": "__end__",
"type": "schema",
"data": "__end__"
}
],
"edges": [
{
"source": "A",
"target": "__end__"
},
{
"source": "B",
"target": "__end__"
},
{
"source": "__start__",
"target": "A"
},
{
"source": "__start__",
"target": "B"
}
]
}
'''
# ---
# name: test_conditional_state_graph_with_list_edge_inputs.1
'''
graph TD;
A --> __end__;
B --> __end__;
__start__ --> A;
__start__ --> B;
'''
# ---
# name: test_conditional_state_graph[postgres]
'{"title": "LangGraphInput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}'
# ---
@@ -3052,6 +3012,76 @@
'''
# ---
# name: test_conditional_state_graph_with_list_edge_inputs
'''
{
"nodes": [
{
"id": "__start__",
"type": "schema",
"data": "__start__"
},
{
"id": "A",
"type": "runnable",
"data": {
"id": [
"langgraph",
"utils",
"RunnableCallable"
],
"name": "A"
}
},
{
"id": "B",
"type": "runnable",
"data": {
"id": [
"langgraph",
"utils",
"RunnableCallable"
],
"name": "B"
}
},
{
"id": "__end__",
"type": "schema",
"data": "__end__"
}
],
"edges": [
{
"source": "A",
"target": "__end__"
},
{
"source": "B",
"target": "__end__"
},
{
"source": "__start__",
"target": "A"
},
{
"source": "__start__",
"target": "B"
}
]
}
'''
# ---
# name: test_conditional_state_graph_with_list_edge_inputs.1
'''
graph TD;
A --> __end__;
B --> __end__;
__start__ --> A;
__start__ --> B;
'''
# ---
# name: test_dynamic_interrupt
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
libs/langgraph/tests/any_str.py
+21 -2
View File
@@ -2,16 +2,35 @@ from typing import Any, Sequence
class AnyStr(str):
def __init__(self) -> None:
def __init__(self, prefix: str = "") -> None:
super().__init__()
self.prefix = prefix
def __eq__(self, other: object) -> bool:
return isinstance(other, str)
return isinstance(other, str) and other.startswith(self.prefix)
def __hash__(self) -> int:
return hash(str(self))
class AnyDict(dict):
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
def __eq__(self, other: object) -> bool:
print("did we get here")
if not isinstance(other, dict) or len(self) != len(other):
return False
for k, v in self.items():
if kk := next((kk for kk in other if kk == k), None):
if v == other[kk]:
continue
else:
return False
else:
return True
class AnyVersion:
def __init__(self) -> None:
super().__init__()
libs/langgraph/tests/test_prebuilt.py
+9 -2
View File
@@ -18,6 +18,7 @@ from langchain_core.tools import BaseTool
from langchain_core.tools import tool as dec_tool
from pydantic import BaseModel as BaseModelV2
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.prebuilt import ToolNode, ValidationNode, create_react_agent
from langgraph.prebuilt.tool_node import InjectedState
from tests.messages import _AnyIdHumanMessage
@@ -55,7 +56,9 @@ class FakeToolCallingModel(BaseChatModel):
["memory", "sqlite", "postgres", "postgres_pipe"],
)
def test_no_modifier(request: pytest.FixtureRequest, checkpointer_name: str) -> None:
checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name)
checkpointer: BaseCheckpointSaver = request.getfixturevalue(
"checkpointer_" + checkpointer_name
)
model = FakeToolCallingModel()
agent = create_react_agent(model, [], checkpointer=checkpointer)
@@ -76,6 +79,7 @@ def test_no_modifier(request: pytest.FixtureRequest, checkpointer_name: str) ->
"agent": "agent",
}
assert saved.metadata == {
"parents": {},
"source": "loop",
"writes": {"agent": {"messages": [AIMessage(content="hi?", id="0")]}},
"step": 1,
@@ -90,7 +94,9 @@ def test_no_modifier(request: pytest.FixtureRequest, checkpointer_name: str) ->
async def test_no_modifier_async(
request: pytest.FixtureRequest, checkpointer_name: str
) -> None:
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
checkpointer: BaseCheckpointSaver = request.getfixturevalue(
f"checkpointer_{checkpointer_name}"
)
model = FakeToolCallingModel()
@@ -112,6 +118,7 @@ async def test_no_modifier_async(
"agent": "agent",
}
assert saved.metadata == {
"parents": {},
"source": "loop",
"writes": {"agent": {"messages": [AIMessage(content="hi?", id="0")]}},
"step": 1,
libs/langgraph/tests/test_pregel.py
+1496 -873
View File
File diff suppressed because it is too large Load Diff
libs/langgraph/tests/test_pregel_async.py
+871 -308
View File
File diff suppressed because it is too large Load Diff