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

add checkpointer=INHERIT_CHECKPOINTER

Browse Source
This commit is contained in:
vbarda
2024-08-12 14:27:50 -04:00
parent 322cfc46d3
commit 9948125745
5 changed files with 568 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/graph/graph.py
+2 -3
View File
@@ -24,7 +24,6 @@ from langchain_core.runnables.graph import Graph as DrawableGraph
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,
@@ -33,7 +32,7 @@ from langgraph.constants import (
Send,
)
from langgraph.errors import InvalidUpdateError
from langgraph.pregel import Channel, Pregel
from langgraph.pregel import Channel, CheckpointerType, Pregel
from langgraph.pregel.read import PregelNode
from langgraph.pregel.types import All
from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
@@ -369,7 +368,7 @@ class Graph:
def compile(
self,
checkpointer: Optional[BaseCheckpointSaver] = None,
checkpointer: Optional[CheckpointerType] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
debug: bool = False,
libs/langgraph/langgraph/graph/state.py
+2 -2
View File
@@ -29,7 +29,6 @@ from langgraph.channels.dynamic_barrier_value import DynamicBarrierValue, WaitFo
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.errors import InvalidUpdateError
from langgraph.graph.graph import (
@@ -41,6 +40,7 @@ from langgraph.graph.graph import (
Send,
)
from langgraph.managed.base import ManagedValue, is_managed_value
from langgraph.pregel import CheckpointerType
from langgraph.pregel.read import ChannelRead, PregelNode
from langgraph.pregel.types import All, RetryPolicy
from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
@@ -373,7 +373,7 @@ class StateGraph(Graph):
def compile(
self,
checkpointer: Optional[BaseCheckpointSaver] = None,
checkpointer: Optional[CheckpointerType] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
debug: bool = False,
libs/langgraph/langgraph/pregel/__init__.py
+21 -2
View File
@@ -13,6 +13,7 @@ from typing import (
Callable,
Dict,
Iterator,
Literal,
Mapping,
Optional,
Sequence,
@@ -115,6 +116,10 @@ WriteValue = Union[
]
INHERIT_CHECKPOINTER = "inherit_checkpointer"
CheckpointerType = Union[BaseCheckpointSaver, Literal["inherit_checkpointer"]]
class Channel:
@overload
@classmethod
@@ -218,7 +223,7 @@ class Pregel(
debug: bool = Field(default_factory=get_debug)
"""Whether to print debug information during execution. Defaults to False."""
checkpointer: Optional[BaseCheckpointSaver] = None
checkpointer: Optional[CheckpointerType] = None
"""Checkpointer used to save and load graph state. Defaults to None."""
retry_policy: Optional[RetryPolicy] = None
@@ -275,6 +280,7 @@ class Pregel(
+ (
self.checkpointer.config_specs
if self.checkpointer is not None
and self.checkpointer != INHERIT_CHECKPOINTER
else []
)
+ (
@@ -950,10 +956,23 @@ 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 self.checkpointer is None:
raise ValueError(
"Missing checkpointer for subgraph. "
"Please compile the subgraph graph with checkpointer=INHERIT_CHECKPOINTER (from langgraph.pregel import INHERIT_CHECKPOINTER)."
)
if self.checkpointer != INHERIT_CHECKPOINTER:
raise ValueError(
"Custom checkpointers for subgraphs are not allowed. "
"Please compile the subgraph graph with checkpointer=INHERIT_CHECKPOINTER (from langgraph.pregel import INHERIT_CHECKPOINTER)."
)
if (
config is not None
and config.get("configurable", {}).get(CONFIG_KEY_CHECKPOINTER)
and (interrupt_after or interrupt_before)
and self.checkpointer == INHERIT_CHECKPOINTER
):
checkpointer: Optional[BaseCheckpointSaver] = config["configurable"][
CONFIG_KEY_CHECKPOINTER
libs/langgraph/tests/test_pregel.py
+272 -8
View File
@@ -60,7 +60,13 @@ from langgraph.prebuilt.chat_agent_executor import (
create_tool_calling_executor,
)
from langgraph.prebuilt.tool_node import ToolNode
from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
from langgraph.pregel import (
INHERIT_CHECKPOINTER,
Channel,
GraphRecursionError,
Pregel,
StateSnapshot,
)
from langgraph.pregel.retry import RetryPolicy
from tests.any_str import AnyStr
from tests.memory_assert import (
@@ -7737,7 +7743,10 @@ def test_nested_graph_interrupts(
graph = StateGraph(State)
graph.add_node("outer_1", outer_1)
graph.add_node("inner", inner.compile(interrupt_before=["inner_2"]))
graph.add_node(
"inner",
inner.compile(interrupt_before=["inner_2"], checkpointer=INHERIT_CHECKPOINTER),
)
graph.add_node("outer_2", outer_2)
graph.set_entry_point("outer_1")
graph.add_edge("outer_1", "inner")
@@ -8913,7 +8922,10 @@ def test_nested_graph_interrupts_parallel(
return {"my_key": " and back again"}
graph = StateGraph(State)
graph.add_node("inner", inner.compile(interrupt_before=["inner_2"]))
graph.add_node(
"inner",
inner.compile(interrupt_before=["inner_2"], checkpointer=INHERIT_CHECKPOINTER),
)
graph.add_node("outer_1", outer_1)
graph.add_node("outer_2", outer_2)
@@ -8997,7 +9009,6 @@ def test_nested_graph_interrupts_parallel(
]
@pytest.mark.skip
@pytest.mark.parametrize(
"checkpointer_name",
["memory", "sqlite", "postgres", "postgres_pipe"],
@@ -9005,7 +9016,7 @@ def test_nested_graph_interrupts_parallel(
def test_doubly_nested_graph_interrupts(
request: pytest.FixtureRequest, checkpointer_name: str
) -> None:
checkpointer = request.getfixturevalue(checkpointer_name)
checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name)
class State(TypedDict):
my_key: str
@@ -9032,7 +9043,12 @@ def test_doubly_nested_graph_interrupts(
grandchild.set_finish_point("grandchild_2")
child = StateGraph(ChildState)
child.add_node("child_1", grandchild.compile(interrupt_before=["grandchild_2"]))
child.add_node(
"child_1",
grandchild.compile(
interrupt_before=["grandchild_2"], checkpointer=INHERIT_CHECKPOINTER
),
)
child.set_entry_point("child_1")
child.set_finish_point("child_1")
@@ -9044,7 +9060,7 @@ def test_doubly_nested_graph_interrupts(
graph = StateGraph(State)
graph.add_node("parent_1", parent_1)
graph.add_node("child", child.compile())
graph.add_node("child", child.compile(checkpointer=INHERIT_CHECKPOINTER))
graph.add_node("parent_2", parent_2)
graph.set_entry_point("parent_1")
graph.add_edge("parent_1", "child")
@@ -9136,7 +9152,10 @@ def test_nested_graph_state(
graph = StateGraph(State)
graph.add_node("outer_1", outer_1)
graph.add_node("inner", inner.compile(interrupt_before=["inner_2"]))
graph.add_node(
"inner",
inner.compile(interrupt_before=["inner_2"], checkpointer=INHERIT_CHECKPOINTER),
)
graph.add_node("outer_2", outer_2)
graph.set_entry_point("outer_1")
graph.add_edge("outer_1", "inner")
@@ -9543,6 +9562,251 @@ def test_nested_graph_state(
]
@pytest.mark.parametrize(
"checkpointer_name",
["memory", "sqlite", "postgres", "postgres_pipe"],
)
def test_doubly_nested_graph_state(
request: pytest.FixtureRequest, checkpointer_name: str
) -> None:
checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name)
class State(TypedDict):
my_key: str
class ChildState(TypedDict):
my_key: str
class GrandChildState(TypedDict):
my_key: str
def grandchild_1(state: ChildState):
return {"my_key": state["my_key"] + " here"}
def grandchild_2(state: ChildState):
return {
"my_key": state["my_key"] + " and there",
}
grandchild = StateGraph(GrandChildState)
grandchild.add_node("grandchild_1", grandchild_1)
grandchild.add_node("grandchild_2", grandchild_2)
grandchild.add_edge("grandchild_1", "grandchild_2")
grandchild.set_entry_point("grandchild_1")
grandchild.set_finish_point("grandchild_2")
child = StateGraph(ChildState)
child.add_node(
"child_1",
grandchild.compile(
interrupt_before=["grandchild_2"], checkpointer=INHERIT_CHECKPOINTER
),
)
child.set_entry_point("child_1")
child.set_finish_point("child_1")
def parent_1(state: State):
return {"my_key": "hi " + state["my_key"]}
def parent_2(state: State):
return {"my_key": state["my_key"] + " and back again"}
graph = StateGraph(State)
graph.add_node("parent_1", parent_1)
graph.add_node("child", child.compile(checkpointer=INHERIT_CHECKPOINTER))
graph.add_node("parent_2", parent_2)
graph.set_entry_point("parent_1")
graph.add_edge("parent_1", "child")
graph.add_edge("child", "parent_2")
graph.set_finish_point("parent_2")
app = graph.compile(checkpointer=checkpointer)
# test invoke w/ nested interrupt
config = {"configurable": {"thread_id": "1"}}
app.invoke({"my_key": "my value"}, config, debug=True)
assert app.get_state(config) == StateSnapshot(
values={"my_key": "hi my value"},
next=("child",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"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(),
}
},
subgraph_state_snapshots=None,
)
assert app.get_state(config, include_subgraph_state=True) == StateSnapshot(
values={"my_key": "hi my value"},
next=("child",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"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(),
}
},
subgraph_state_snapshots={
"child": StateSnapshot(
values={"my_key": "hi my value"},
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child",
"checkpoint_id": AnyStr(),
}
},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"child_1": StateSnapshot(
values={"my_key": "hi my value here"},
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child|child_1",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"grandchild_1": {"my_key": "hi my value here"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child|child_1",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
)
},
)
app.invoke(None, config, debug=True)
assert app.get_state(config, include_subgraph_state=True) == StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"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(),
}
},
subgraph_state_snapshots={
"child": StateSnapshot(
values={"my_key": "hi my value here and there"},
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"child_1": {"my_key": "hi my value here and there"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"child_1": StateSnapshot(
values={"my_key": "hi my value here and there"},
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child|child_1",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {
"grandchild_2": {"my_key": "hi my value here and there"}
},
"step": 2,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child|child_1",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
)
},
)
def test_repeat_condition(snapshot: SnapshotAssertion) -> None:
class AgentState(TypedDict):
hello: str
libs/langgraph/tests/test_pregel_async.py
+271 -7
View File
@@ -56,7 +56,13 @@ from langgraph.prebuilt.chat_agent_executor import (
)
from langgraph.prebuilt.tool_executor import ToolExecutor
from langgraph.prebuilt.tool_node import ToolNode
from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
from langgraph.pregel import (
INHERIT_CHECKPOINTER,
Channel,
GraphRecursionError,
Pregel,
StateSnapshot,
)
from langgraph.pregel.retry import RetryPolicy
from tests.any_str import AnyStr
from tests.memory_assert import (
@@ -6232,7 +6238,10 @@ async def test_nested_graph_interrupts(
graph = StateGraph(State)
graph.add_node("outer_1", outer_1)
graph.add_node("inner", inner.compile(interrupt_before=["inner_2"]))
graph.add_node(
"inner",
inner.compile(interrupt_before=["inner_2"], checkpointer=INHERIT_CHECKPOINTER),
)
graph.add_node("outer_2", outer_2)
graph.set_entry_point("outer_1")
graph.add_edge("outer_1", "inner")
@@ -7415,7 +7424,10 @@ async def test_nested_graph_interrupts_parallel(
return {"my_key": " and back again"}
graph = StateGraph(State)
graph.add_node("inner", inner.compile(interrupt_before=["inner_2"]))
graph.add_node(
"inner",
inner.compile(interrupt_before=["inner_2"], checkpointer=INHERIT_CHECKPOINTER),
)
graph.add_node("outer_1", outer_1)
graph.add_node("outer_2", outer_2)
@@ -7501,7 +7513,6 @@ async def test_nested_graph_interrupts_parallel(
]
@pytest.mark.skip
@pytest.mark.parametrize(
"checkpointer_name",
["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"],
@@ -7536,7 +7547,12 @@ async def test_doubly_nested_graph_interrupts(
grandchild.set_finish_point("grandchild_2")
child = StateGraph(ChildState)
child.add_node("child_1", grandchild.compile(interrupt_before=["grandchild_2"]))
child.add_node(
"child_1",
grandchild.compile(
interrupt_before=["grandchild_2"], checkpointer=INHERIT_CHECKPOINTER
),
)
child.set_entry_point("child_1")
child.set_finish_point("child_1")
@@ -7548,7 +7564,7 @@ async def test_doubly_nested_graph_interrupts(
graph = StateGraph(State)
graph.add_node("parent_1", parent_1)
graph.add_node("child", child.compile())
graph.add_node("child", child.compile(checkpointer=INHERIT_CHECKPOINTER))
graph.add_node("parent_2", parent_2)
graph.set_entry_point("parent_1")
graph.add_edge("parent_1", "child")
@@ -7643,7 +7659,10 @@ async def test_nested_graph_state(
graph = StateGraph(State)
graph.add_node("outer_1", outer_1)
graph.add_node("inner", inner.compile(interrupt_before=["inner_2"]))
graph.add_node(
"inner",
inner.compile(interrupt_before=["inner_2"], checkpointer=INHERIT_CHECKPOINTER),
)
graph.add_node("outer_2", outer_2)
graph.set_entry_point("outer_1")
graph.add_edge("outer_1", "inner")
@@ -8052,6 +8071,251 @@ async def test_nested_graph_state(
]
@pytest.mark.parametrize(
"checkpointer_name",
["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"],
)
async def test_doubly_nested_graph_state(
request: pytest.FixtureRequest, checkpointer_name: str
) -> None:
checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name)
class State(TypedDict):
my_key: str
class ChildState(TypedDict):
my_key: str
class GrandChildState(TypedDict):
my_key: str
def grandchild_1(state: ChildState):
return {"my_key": state["my_key"] + " here"}
def grandchild_2(state: ChildState):
return {
"my_key": state["my_key"] + " and there",
}
grandchild = StateGraph(GrandChildState)
grandchild.add_node("grandchild_1", grandchild_1)
grandchild.add_node("grandchild_2", grandchild_2)
grandchild.add_edge("grandchild_1", "grandchild_2")
grandchild.set_entry_point("grandchild_1")
grandchild.set_finish_point("grandchild_2")
child = StateGraph(ChildState)
child.add_node(
"child_1",
grandchild.compile(
interrupt_before=["grandchild_2"], checkpointer=INHERIT_CHECKPOINTER
),
)
child.set_entry_point("child_1")
child.set_finish_point("child_1")
def parent_1(state: State):
return {"my_key": "hi " + state["my_key"]}
def parent_2(state: State):
return {"my_key": state["my_key"] + " and back again"}
graph = StateGraph(State)
graph.add_node("parent_1", parent_1)
graph.add_node("child", child.compile(checkpointer=INHERIT_CHECKPOINTER))
graph.add_node("parent_2", parent_2)
graph.set_entry_point("parent_1")
graph.add_edge("parent_1", "child")
graph.add_edge("child", "parent_2")
graph.set_finish_point("parent_2")
app = graph.compile(checkpointer=checkpointer)
# test invoke w/ nested interrupt
config = {"configurable": {"thread_id": "1"}}
await app.ainvoke({"my_key": "my value"}, config, debug=True)
assert await app.aget_state(config) == StateSnapshot(
values={"my_key": "hi my value"},
next=("child",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"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(),
}
},
subgraph_state_snapshots=None,
)
assert await app.aget_state(config, include_subgraph_state=True) == StateSnapshot(
values={"my_key": "hi my value"},
next=("child",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"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(),
}
},
subgraph_state_snapshots={
"child": StateSnapshot(
values={"my_key": "hi my value"},
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child",
"checkpoint_id": AnyStr(),
}
},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"child_1": StateSnapshot(
values={"my_key": "hi my value here"},
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child|child_1",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"grandchild_1": {"my_key": "hi my value here"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child|child_1",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
)
},
)
await app.ainvoke(None, config, debug=True)
assert await app.aget_state(config, include_subgraph_state=True) == StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"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(),
}
},
subgraph_state_snapshots={
"child": StateSnapshot(
values={"my_key": "hi my value here and there"},
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"child_1": {"my_key": "hi my value here and there"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots={
"child_1": StateSnapshot(
values={"my_key": "hi my value here and there"},
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child|child_1",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {
"grandchild_2": {"my_key": "hi my value here and there"}
},
"step": 2,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "child|child_1",
"checkpoint_id": AnyStr(),
}
},
subgraph_state_snapshots=None,
)
},
)
},
)
async def test_checkpoint_metadata() -> None:
"""This test verifies that a run's configurable fields are merged with the
previous checkpoint config for each step in the run.