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

Merge pull request #1108 from langchain-ai/vb/update-get-state

Browse Source 
langgraph: update get_state to handle nested subgraph state
This commit is contained in:
Nuno Campos
2024-08-26 17:58:41 -07:00
committed by GitHub
parent ce0fad8ddb 507930e5b7
commit 65784d8174
19 changed files with 3324 additions and 317 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py
+7 -3
View File
@@ -255,9 +255,13 @@ class BasePostgresSaver(BaseCheckpointSaver):
if config:
wheres.append("thread_id = %s ")
param_values.append(config["configurable"]["thread_id"])
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
wheres.append("checkpoint_ns = %s")
param_values.append(checkpoint_ns)
if checkpoint_ns := config["configurable"].get("checkpoint_ns"):
wheres.append("checkpoint_ns = %s")
param_values.append(checkpoint_ns)
if checkpoint_id := get_checkpoint_id(config):
wheres.append("checkpoint_id = %s ")
param_values.append(checkpoint_id)
# construct predicate for metadata filter
if filter:
libs/checkpoint-postgres/tests/test_async.py
+6 -20
View File
@@ -87,29 +87,15 @@ class TestAsyncPostgresSaver:
search_results_4 = [c async for c in saver.alist(None, filter=query_4)]
assert len(search_results_4) == 0
# search by config (defaults to root graph checkpoints)
# search by config (defaults to checkpoints across all namespaces)
search_results_5 = [
c
async for c in saver.alist({"configurable": {"thread_id": "thread-2"}})
]
assert len(search_results_5) == 1
assert search_results_5[0].config["configurable"]["checkpoint_ns"] == ""
# search by config and checkpoint_ns
search_results_6 = [
c
async for c in saver.alist(
{
"configurable": {
"thread_id": "thread-2",
"checkpoint_ns": "inner",
}
}
)
]
assert len(search_results_6) == 1
assert (
search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner"
)
assert len(search_results_5) == 2
assert {
search_results_5[0].config["configurable"]["checkpoint_ns"],
search_results_5[1].config["configurable"]["checkpoint_ns"],
} == {"", "inner"}
# TODO: test before and limit params
libs/checkpoint-postgres/tests/test_sync.py
+6 -19
View File
@@ -88,27 +88,14 @@ class TestPostgresSaver:
search_results_4 = list(saver.list(None, filter=query_4))
assert len(search_results_4) == 0
# search by config (defaults to root graph checkpoints)
# search by config (defaults to checkpoints across all namespaces)
search_results_5 = list(
saver.list({"configurable": {"thread_id": "thread-2"}})
)
assert len(search_results_5) == 1
assert search_results_5[0].config["configurable"]["checkpoint_ns"] == ""
# search by config and checkpoint_ns
search_results_6 = list(
saver.list(
{
"configurable": {
"thread_id": "thread-2",
"checkpoint_ns": "inner",
}
}
)
)
assert len(search_results_6) == 1
assert (
search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner"
)
assert len(search_results_5) == 2
assert {
search_results_5[0].config["configurable"]["checkpoint_ns"],
search_results_5[1].config["configurable"]["checkpoint_ns"],
} == {"", "inner"}
# TODO: test before and limit params
libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/utils.py
+7 -3
View File
@@ -70,9 +70,13 @@ def search_where(
if config is not None:
wheres.append("thread_id = ?")
param_values.append(config["configurable"]["thread_id"])
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
wheres.append("checkpoint_ns = ?")
param_values.append(checkpoint_ns)
if checkpoint_ns := config["configurable"].get("checkpoint_ns"):
wheres.append("checkpoint_ns = ?")
param_values.append(checkpoint_ns)
if checkpoint_id := get_checkpoint_id(config):
wheres.append("checkpoint_id = ?")
param_values.append(checkpoint_id)
# construct predicate for metadata filter
if filter:
libs/checkpoint-sqlite/tests/test_aiosqlite.py
+6 -20
View File
@@ -84,29 +84,15 @@ class TestAsyncSqliteSaver:
search_results_4 = [c async for c in saver.alist(None, filter=query_4)]
assert len(search_results_4) == 0
# search by config (defaults to root graph checkpoints)
# search by config (defaults to checkpoints across all namespaces)
search_results_5 = [
c
async for c in saver.alist({"configurable": {"thread_id": "thread-2"}})
]
assert len(search_results_5) == 1
assert search_results_5[0].config["configurable"]["checkpoint_ns"] == ""
# search by config and checkpoint_ns
search_results_6 = [
c
async for c in saver.alist(
{
"configurable": {
"thread_id": "thread-2",
"checkpoint_ns": "inner",
}
}
)
]
assert len(search_results_6) == 1
assert (
search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner"
)
assert len(search_results_5) == 2
assert {
search_results_5[0].config["configurable"]["checkpoint_ns"],
search_results_5[1].config["configurable"]["checkpoint_ns"],
} == {"", "inner"}
# TODO: test before and limit params
libs/checkpoint-sqlite/tests/test_sqlite.py
+6 -19
View File
@@ -87,28 +87,15 @@ class TestSqliteSaver:
search_results_4 = list(saver.list(None, filter=query_4))
assert len(search_results_4) == 0
# search by config (defaults to root graph checkpoints)
# search by config (defaults to checkpoints across all namespaces)
search_results_5 = list(
saver.list({"configurable": {"thread_id": "thread-2"}})
)
assert len(search_results_5) == 1
assert search_results_5[0].config["configurable"]["checkpoint_ns"] == ""
# search by config and checkpoint_ns
search_results_6 = list(
saver.list(
{
"configurable": {
"thread_id": "thread-2",
"checkpoint_ns": "inner",
}
}
)
)
assert len(search_results_6) == 1
assert (
search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner"
)
assert len(search_results_5) == 2
assert {
search_results_5[0].config["configurable"]["checkpoint_ns"],
search_results_5[1].config["configurable"]["checkpoint_ns"],
} == {"", "inner"}
# TODO: test before and limit params
libs/checkpoint/langgraph/checkpoint/memory/__init__.py
+70 -48
View File
@@ -177,62 +177,77 @@ class MemorySaver(
Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples.
"""
thread_ids = (config["configurable"]["thread_id"],) if config else self.storage
checkpoint_ns = (
config["configurable"].get("checkpoint_ns", "") if config else ""
config_checkpoint_ns = (
config["configurable"].get("checkpoint_ns") if config else None
)
config_checkpoint_id = get_checkpoint_id(config) if config else None
for thread_id in thread_ids:
for checkpoint_id, (checkpoint, metadata_b, parent_checkpoint_id) in sorted(
self.storage[thread_id][checkpoint_ns].items(),
key=lambda x: x[0],
reverse=True,
):
# filter by checkpoint ID
if (
before
and (before_checkpoint_id := get_checkpoint_id(before))
and checkpoint_id >= before_checkpoint_id
):
for checkpoint_ns in self.storage[thread_id].keys():
if config_checkpoint_ns and checkpoint_ns != config_checkpoint_ns:
continue
# filter by metadata
metadata = self.serde.loads_typed(metadata_b)
if filter and not all(
query_value == metadata[query_key]
for query_key, query_value in filter.items()
for checkpoint_id, (
checkpoint,
metadata_b,
parent_checkpoint_id,
) in sorted(
self.storage[thread_id][checkpoint_ns].items(),
key=lambda x: x[0],
reverse=True,
):
continue
# filter by checkpoint ID from config
if config_checkpoint_id and checkpoint_id != config_checkpoint_id:
continue
# limit search results
if limit is not None and limit <= 0:
break
elif limit is not None:
limit -= 1
# filter by checkpoint ID from `before` config
if (
before
and (before_checkpoint_id := get_checkpoint_id(before))
and checkpoint_id >= before_checkpoint_id
):
continue
writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values()
# filter by metadata
metadata = self.serde.loads_typed(metadata_b)
if filter and not all(
query_value == metadata.get(query_key)
for query_key, query_value in filter.items()
):
continue
yield CheckpointTuple(
config={
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint_id,
# limit search results
if limit is not None and limit <= 0:
break
elif limit is not None:
limit -= 1
writes = self.writes[
(thread_id, checkpoint_ns, checkpoint_id)
].values()
yield CheckpointTuple(
config={
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": checkpoint_id,
}
},
checkpoint=self.serde.loads_typed(checkpoint),
metadata=metadata,
parent_config={
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": parent_checkpoint_id,
}
}
},
checkpoint=self.serde.loads_typed(checkpoint),
metadata=metadata,
parent_config={
"configurable": {
"thread_id": thread_id,
"checkpoint_ns": checkpoint_ns,
"checkpoint_id": parent_checkpoint_id,
}
}
if parent_checkpoint_id
else None,
pending_writes=[
(id, c, self.serde.loads_typed(v)) for id, c, v in writes
],
)
if parent_checkpoint_id
else None,
pending_writes=[
(id, c, self.serde.loads_typed(v)) for id, c, v in writes
],
)
def put(
self,
@@ -338,7 +353,14 @@ class MemorySaver(
"""
loop = asyncio.get_running_loop()
iter = await loop.run_in_executor(
None, partial(self.list, before=before, limit=limit, filter=filter), config
None,
partial(
self.list,
before=before,
limit=limit,
filter=filter,
),
config,
)
while True:
# handling StopIteration exception inside coroutine won't work
libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py
+1 -1
View File
@@ -109,7 +109,7 @@ class JsonPlusSerializer(SerializerProtocol):
return self._encode_constructor_args(obj.__class__, args=[obj.value])
elif isinstance(obj, SendProtocol):
return self._encode_constructor_args(
obj.__class__, kwargs={"node": obj.node, "arg": obj.arg}
obj.__class__, kwargs={"node": obj.node, "arg": obj.arg, "id": obj.id}
)
elif isinstance(obj, (bytes, bytearray)):
return self._encode_constructor_args(
libs/checkpoint/langgraph/checkpoint/serde/types.py
+1
View File
@@ -57,6 +57,7 @@ class SendProtocol(Protocol):
# Mirrors langgraph.constants.Send
node: str
arg: Any
id: str
def __hash__(self) -> int:
...
libs/checkpoint/tests/test_memory.py
+9 -14
View File
@@ -82,26 +82,20 @@ class TestMemorySaver:
assert search_results_2[0].metadata == self.metadata_2
search_results_3 = list(self.memory_saver.list(None, filter=query_3))
assert len(search_results_3) == 2
assert len(search_results_3) == 3
search_results_4 = list(self.memory_saver.list(None, filter=query_4))
assert len(search_results_4) == 0
# search by config (defaults to root graph checkpoints)
# search by config (defaults to checkpoints across all namespaces)
search_results_5 = list(
self.memory_saver.list({"configurable": {"thread_id": "thread-2"}})
)
assert len(search_results_5) == 1
assert search_results_5[0].config["configurable"]["checkpoint_ns"] == ""
# search by config and checkpoint_ns
search_results_6 = list(
self.memory_saver.list(
{"configurable": {"thread_id": "thread-2", "checkpoint_ns": "inner"}}
)
)
assert len(search_results_6) == 1
assert search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner"
assert len(search_results_5) == 2
assert {
search_results_5[0].config["configurable"]["checkpoint_ns"],
search_results_5[1].config["configurable"]["checkpoint_ns"],
} == {"", "inner"}
# TODO: test before and limit params
@@ -110,6 +104,7 @@ class TestMemorySaver:
# save checkpoints
self.memory_saver.put(self.config_1, self.chkpnt_1, self.metadata_1, {})
self.memory_saver.put(self.config_2, self.chkpnt_2, self.metadata_2, {})
self.memory_saver.put(self.config_3, self.chkpnt_3, self.metadata_3, {})
# call method / assertions
query_1: CheckpointMetadata = {"source": "input"} # search by 1 key
@@ -135,7 +130,7 @@ class TestMemorySaver:
search_results_3 = [
c async for c in self.memory_saver.alist(None, filter=query_3)
]
assert len(search_results_3) == 2
assert len(search_results_3) == 3
search_results_4 = [
c async for c in self.memory_saver.alist(None, filter=query_4)
libs/langgraph/langgraph/constants.py
+10 -4
View File
@@ -1,5 +1,6 @@
from dataclasses import dataclass
from typing import Any, Literal
from typing import Any, Literal, Optional
from uuid import uuid4
INPUT = "__input__"
CONFIG_KEY_SEND = "__pregel_send"
@@ -31,6 +32,7 @@ START = "__start__"
END = "__end__"
CHECKPOINT_NAMESPACE_SEPARATOR = "|"
SEND_CHECKPOINT_NAMESPACE_SEPARATOR = ":"
class Send:
@@ -49,6 +51,7 @@ class Send:
Attributes:
node (str): The name of the target node to send the message to.
arg (Any): The state or message to send to the target node.
id (str): ID associated with the Send.
Examples:
>>> from typing import Annotated
@@ -76,23 +79,26 @@ class Send:
node: str
arg: Any
id: Optional[str]
def __init__(self, /, node: str, arg: Any) -> None:
def __init__(self, /, node: str, arg: Any, id: Optional[str] = None) -> None:
"""
Initialize a new instance of the Send class.
Args:
node (str): The name of the target node to send the message to.
arg (Any): The state or message to send to the target node.
id (str): ID associated with the Send.
"""
self.node = node
self.arg = arg
self.id = id or str(uuid4())
def __hash__(self) -> int:
return hash((self.node, self.arg))
return hash((self.node, self.arg, self.id))
def __repr__(self) -> str:
return f"Send(node={self.node!r}, arg={self.arg!r})"
return f"Send(node={self.node!r}, arg={self.arg!r}, id={self.id!r})"
def __eq__(self, value: object) -> bool:
return (
libs/langgraph/langgraph/graph/graph.py
+10 -4
View File
@@ -28,6 +28,7 @@ from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.constants import (
CHECKPOINT_NAMESPACE_SEPARATOR,
END,
SEND_CHECKPOINT_NAMESPACE_SEPARATOR,
START,
TAG_HIDDEN,
Send,
@@ -160,10 +161,15 @@ class Graph:
*,
metadata: Optional[dict[str, Any]] = None,
) -> None:
if isinstance(node, str) and CHECKPOINT_NAMESPACE_SEPARATOR in node:
raise ValueError(
f"'{CHECKPOINT_NAMESPACE_SEPARATOR}' is a reserved character and is not allowed in the node names."
)
if isinstance(node, str):
for character in (
CHECKPOINT_NAMESPACE_SEPARATOR,
SEND_CHECKPOINT_NAMESPACE_SEPARATOR,
):
if character in node:
raise ValueError(
f"'{character}' is a reserved character and is not allowed in the node names."
)
if self.compiled:
logger.warning(
libs/langgraph/langgraph/graph/state.py
+13 -5
View File
@@ -31,7 +31,11 @@ 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, TAG_HIDDEN
from langgraph.constants import (
CHECKPOINT_NAMESPACE_SEPARATOR,
SEND_CHECKPOINT_NAMESPACE_SEPARATOR,
TAG_HIDDEN,
)
from langgraph.errors import InvalidUpdateError
from langgraph.graph.graph import (
END,
@@ -317,10 +321,14 @@ class StateGraph(Graph):
if node == END or node == START:
raise ValueError(f"Node `{node}` is reserved.")
if CHECKPOINT_NAMESPACE_SEPARATOR in node:
raise ValueError(
f"'{CHECKPOINT_NAMESPACE_SEPARATOR}' is a reserved character and is not allowed in the node names."
)
for character in (
CHECKPOINT_NAMESPACE_SEPARATOR,
SEND_CHECKPOINT_NAMESPACE_SEPARATOR,
):
if character in node:
raise ValueError(
f"'{character}' is a reserved character and is not allowed in the node names."
)
try:
if isfunction(action) and (
libs/langgraph/langgraph/pregel/__init__.py
+278 -123
View File
@@ -54,17 +54,20 @@ from langgraph.channels.base import (
)
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
CheckpointTuple,
copy_checkpoint,
create_checkpoint,
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,
)
from langgraph.errors import GraphInterrupt, GraphRecursionError, InvalidUpdateError
from langgraph.managed.base import ManagedValueSpec
@@ -80,6 +83,7 @@ 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
@@ -91,7 +95,9 @@ from langgraph.pregel.types import (
StateSnapshot,
StreamMode,
)
from langgraph.pregel.utils import get_new_channel_versions
from langgraph.pregel.utils import (
get_new_channel_versions,
)
from langgraph.pregel.validate import validate_graph, validate_keys
from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
from langgraph.store.base import BaseStore
@@ -178,6 +184,181 @@ 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]]
):
@@ -347,6 +528,18 @@ 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():
if isinstance(node.bound, Pregel):
yield node.bound
yield from node.bound.subgraphs
elif isinstance(node.bound, RunnableSequence):
for runnable in node.bound.steps:
if isinstance(runnable, Pregel):
yield runnable
yield from runnable.subgraphs
def get_state(self, config: RunnableConfig) -> StateSnapshot:
"""Get the current state of the graph."""
if not self.checkpointer:
@@ -354,31 +547,18 @@ class Pregel(
config = merge_configs(self.config, config) if self.config else config
saved = self.checkpointer.get_tuple(config)
checkpoint = saved.checkpoint if saved else empty_checkpoint()
config = saved.config if saved else config
with ChannelsManager(self.channels, checkpoint, config, skip_context=True) as (
channels,
managed,
):
next_tasks = prepare_next_tasks(
checkpoint,
self.nodes,
channels,
managed,
config,
saved.metadata.get("step", -1) + 1 if saved else -1,
for_execution=False,
)
checkpoint_config = saved.config if saved else config
checkpoint_ns_to_graph: dict[str, Pregel] = _get_checkpoint_ns_to_graph(self)
return StateSnapshot(
read_channels(channels, self.stream_channels_asis),
tuple(t.name for t in next_tasks),
saved.config if saved else config,
saved.metadata if saved else None,
saved.checkpoint["ts"] if saved else None,
saved.parent_config if saved else None,
tasks_w_writes(next_tasks, saved.pending_writes if saved else None),
)
# 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
)
async def aget_state(self, config: RunnableConfig) -> StateSnapshot:
"""Get the current state of the graph."""
@@ -387,30 +567,22 @@ class Pregel(
config = merge_configs(self.config, config) if self.config else config
saved = await self.checkpointer.aget_tuple(config)
checkpoint = saved.checkpoint if saved else empty_checkpoint()
checkpoint_config = saved.config if saved else config
checkpoint_ns_to_graph: dict[str, Pregel] = _get_checkpoint_ns_to_graph(self)
config = saved.config if saved else config
async with AsyncChannelsManager(
self.channels, checkpoint, config, skip_context=True
) as (channels, managed):
next_tasks = prepare_next_tasks(
checkpoint,
self.nodes,
channels,
managed,
config,
saved.metadata.get("step", -1) + 1 if saved else -1,
for_execution=False,
)
return StateSnapshot(
read_channels(channels, self.stream_channels_asis),
tuple(t.name for t in next_tasks),
saved.config if saved else config,
saved.metadata if saved else None,
saved.checkpoint["ts"] if saved else None,
saved.parent_config if saved else None,
tasks_w_writes(next_tasks, saved.pending_writes if saved else None),
)
# 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
)
def get_state_history(
self,
@@ -428,39 +600,30 @@ class Pregel(
and signature(self.checkpointer.list).parameters.get("filter") is None
):
raise ValueError("Checkpointer does not support filtering")
for (
config,
checkpoint,
metadata,
parent_config,
pending_writes,
) in self.checkpointer.list(
merge_configs(self.config, config) if self.config else config,
before=before,
limit=limit,
filter=filter,
):
with ChannelsManager(
self.channels, checkpoint, config, skip_context=True
) as (channels, managed):
next_tasks = prepare_next_tasks(
checkpoint,
self.nodes,
channels,
managed,
config,
metadata.get("step", -1) + 1,
for_execution=False,
)
yield StateSnapshot(
read_channels(channels, self.stream_channels_asis),
tuple(t.name for t in next_tasks),
config,
metadata,
checkpoint["ts"],
parent_config,
tasks_w_writes(next_tasks, pending_writes),
)
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,
)
]
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,
@@ -478,39 +641,36 @@ class Pregel(
and signature(self.checkpointer.list).parameters.get("filter") is None
):
raise ValueError("Checkpointer does not support filtering")
async for (
config,
checkpoint,
metadata,
parent_config,
pending_writes,
) in self.checkpointer.alist(
merge_configs(self.config, config) if self.config else config,
before=before,
limit=limit,
filter=filter,
):
async with AsyncChannelsManager(
self.channels, checkpoint, config, skip_context=True
) as (channels, managed):
next_tasks = prepare_next_tasks(
checkpoint,
self.nodes,
channels,
managed,
config,
metadata.get("step", -1) + 1,
for_execution=False,
)
yield StateSnapshot(
read_channels(channels, self.stream_channels_asis),
tuple(t.name for t in next_tasks),
config,
metadata,
checkpoint["ts"],
parent_config,
tasks_w_writes(next_tasks, pending_writes),
)
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,
)
]
# 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()
)
yield state_snapshot
def update_state(
self,
@@ -816,7 +976,7 @@ class Pregel(
if (
config is not None
and config.get("configurable", {}).get(CONFIG_KEY_CHECKPOINTER)
and (interrupt_after or interrupt_before)
and (interrupt_after or interrupt_before or _has_nested_interrupts(self))
):
checkpointer: Optional[BaseCheckpointSaver] = config["configurable"][
CONFIG_KEY_CHECKPOINTER
@@ -1023,6 +1183,7 @@ class Pregel(
)
else:
loop.put_writes(task.id, [(ERROR, exc)])
else:
# save task writes to checkpointer
loop.put_writes(task.id, task.writes)
@@ -1263,6 +1424,7 @@ class Pregel(
)
if not done:
break # timed out
for fut in done:
task = futures.pop(fut)
if exc := _exception(fut):
@@ -1273,6 +1435,7 @@ class Pregel(
)
else:
loop.put_writes(task.id, [(ERROR, exc)])
else:
# save task writes to checkpointer
loop.put_writes(task.id, task.writes)
@@ -1480,11 +1643,3 @@ def _panic_or_proceed(
inflight.pop().cancel()
# raise timeout error
raise timeout_exc_cls(f"Timed out at step {step}")
def _with_mode(mode: StreamMode, on: bool, iter: Iterator[Any]) -> Iterator[Any]:
if on:
for chunk in iter:
yield (mode, chunk)
else:
yield from iter
libs/langgraph/langgraph/pregel/algo.py
+6 -4
View File
@@ -293,9 +293,9 @@ def prepare_next_tasks(
"langgraph_task_idx": len(tasks),
}
checkpoint_ns = (
f"{parent_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{packet.node}"
f"{parent_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{packet.node}:{packet.id}"
if parent_ns
else packet.node
else f"{packet.node}:{packet.id}"
)
task_id = str(
uuid5(UUID(checkpoint["id"]), json.dumps((checkpoint_ns, metadata)))
@@ -343,8 +343,10 @@ def prepare_next_tasks(
PregelTaskWrites(packet.node, writes, triggers),
config,
),
# in Send we can't checkpoint nested graphs
# as they could be running in parallel
CONFIG_KEY_CHECKPOINTER: checkpointer,
CONFIG_KEY_RESUMING: is_resuming,
"checkpoint_id": checkpoint["id"],
"checkpoint_ns": checkpoint_ns,
},
),
triggers,
libs/langgraph/langgraph/pregel/get_state.py
+36
View File
@@ -0,0 +1,36 @@
from langgraph.constants import CHECKPOINT_NAMESPACE_SEPARATOR
from langgraph.pregel.types import StateSnapshot
def assemble_state_snapshot_hierarchy(
root_checkpoint_ns: str,
checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot],
) -> StateSnapshot:
checkpoint_ns_list_to_visit = sorted(
checkpoint_ns_to_state_snapshots.keys(),
key=lambda x: len(x.split(CHECKPOINT_NAMESPACE_SEPARATOR)),
)
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)
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 {}),
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
)
state_snapshot = checkpoint_ns_to_state_snapshots.pop(root_checkpoint_ns, None)
if state_snapshot is None:
raise ValueError(f"Missing checkpoint for checkpoint NS '{root_checkpoint_ns}'")
return state_snapshot
libs/langgraph/langgraph/pregel/types.py
+2
View File
@@ -92,6 +92,8 @@ 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/test_pregel.py
+1415 -14
View File
File diff suppressed because it is too large Load Diff
libs/langgraph/tests/test_pregel_async.py
+1435 -16
View File
File diff suppressed because it is too large Load Diff