Remove send v2

This commit is contained in:
Nuno Campos
2025-01-14 15:49:57 -08:00
parent 400d83708a
commit 0b74e25f72
13 changed files with 576 additions and 1313 deletions
-3
View File
@@ -1,5 +1,4 @@
import sys
from os import getenv
from types import MappingProxyType
from typing import Any, Literal, Mapping, cast
@@ -93,8 +92,6 @@ NS_END = sys.intern(":")
# for checkpoint_ns, for each level, separates the namespace from the task_id
CONF = cast(Literal["configurable"], sys.intern("configurable"))
# key for the configurable dict in RunnableConfig
FF_SEND_V2 = getenv("LANGGRAPH_FF_SEND_V2", "false").lower() == "true"
# temporary flag to enable new Send semantics
NULL_TASK_ID = sys.intern("00000000-0000-0000-0000-000000000000")
# the task_id to use for writes that are not associated with a task
+4 -2
View File
@@ -966,7 +966,7 @@ class Pregel(PregelProtocol):
return patch_checkpoint_map(
next_config, saved.metadata if saved else None
)
# no values, copy checkpoint
# no values, empty checkpoint
if values is None and as_node is None:
next_checkpoint = create_checkpoint(checkpoint, None, step)
# copy checkpoint
@@ -985,6 +985,7 @@ class Pregel(PregelProtocol):
return patch_checkpoint_map(
next_config, saved.metadata if saved else None
)
# no values, copy checkpoint
if values is None and as_node == "__copy__":
next_checkpoint = create_checkpoint(checkpoint, None, step)
# copy checkpoint
@@ -1248,7 +1249,7 @@ class Pregel(PregelProtocol):
return patch_checkpoint_map(
next_config, saved.metadata if saved else None
)
# no values, copy checkpoint
# no values, empty checkpoint
if values is None and as_node is None:
next_checkpoint = create_checkpoint(checkpoint, None, step)
# copy checkpoint
@@ -1267,6 +1268,7 @@ class Pregel(PregelProtocol):
return patch_checkpoint_map(
next_config, saved.metadata if saved else None
)
# no values, copy checkpoint
if values is None and as_node == "__copy__":
next_checkpoint = create_checkpoint(checkpoint, None, step)
# copy checkpoint
+10 -103
View File
@@ -228,7 +228,7 @@ def apply_writes(
# sort tasks on path, to ensure deterministic order for update application
# any path parts after the 3rd are ignored for sorting
# (we use them for eg. task ids which aren't good for sorting)
tasks = sorted(tasks, key=lambda t: t.path[:3])
tasks = sorted(tasks, key=lambda t: _tuple_str(t.path[:3]))
# if no task has triggers this is applying writes from the null task only
# so we don't do anything other than update the channels written to
bump_step = any(t.triggers for t in tasks)
@@ -273,7 +273,7 @@ def apply_writes(
for chan, val in task.writes:
if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT, RETURN, ERROR):
pass
elif chan == TASKS: # TODO: remove branch in 1.0
elif chan == TASKS:
checkpoint["pending_sends"].append(val)
elif chan in channels:
pending_writes_by_channel[chan].append(val)
@@ -363,8 +363,8 @@ def prepare_next_tasks(
This is the union of all PUSH tasks (Sends) and PULL tasks (nodes triggered
by edges)."""
tasks: list[Union[PregelTask, PregelExecutableTask]] = []
# Consume pending_sends from previous step (legacy version of Send)
for idx, _ in enumerate(checkpoint["pending_sends"]): # TODO: remove branch in 1.0
# Consume pending_sends from previous step
for idx, _ in enumerate(checkpoint["pending_sends"]):
if task := prepare_single_task(
(PUSH, idx),
None,
@@ -400,65 +400,7 @@ def prepare_next_tasks(
manager=manager,
):
tasks.append(task)
# Consume pending Sends from this step (new version of Send)
if any(c == PUSH for _, c, _ in pending_writes):
# group writes by task id
grouped_by_task = defaultdict(list)
for tid, c, _ in pending_writes:
grouped_by_task[tid].append(c)
# prepare send tasks from grouped writes
# 1. start from sends originating from existing tasks
tidx = 0
while tidx < len(tasks):
task = tasks[tidx]
if twrites := grouped_by_task.pop(task.id, None):
for idx, c in enumerate(twrites):
if c != PUSH:
continue
if next_task := prepare_single_task(
(PUSH, task.path, idx, task.id),
None,
checkpoint=checkpoint,
pending_writes=pending_writes,
processes=processes,
channels=channels,
managed=managed,
config=config,
step=step,
for_execution=for_execution,
store=store,
checkpointer=checkpointer,
manager=manager,
):
tasks.append(next_task)
tidx += 1
# key tasks by id
task_map = {t.id: t for t in tasks}
# 2. create new tasks for remaining sends (eg. from update_state)
for tid, writes in grouped_by_task.items():
task = task_map.get(tid)
for idx, c in enumerate(writes):
if c != PUSH:
continue
if next_task := prepare_single_task(
(PUSH, task.path if task else (), idx, tid),
None,
checkpoint=checkpoint,
pending_writes=pending_writes,
processes=processes,
channels=channels,
managed=managed,
config=config,
step=step,
for_execution=for_execution,
store=store,
checkpointer=checkpointer,
manager=manager,
):
task_map[next_task.id] = next_task
else:
task_map = {t.id: t for t in tasks}
return task_map
return {t.id: t for t in tasks}
def prepare_single_task(
@@ -571,8 +513,8 @@ def prepare_single_task(
else:
return PregelTask(task_id, name, task_path[:3])
elif task_path[0] == PUSH:
if len(task_path) == 2: # TODO: remove branch in 1.0
# legacy SEND tasks, executed in superstep n+1
if len(task_path) == 2:
# SEND tasks, executed in superstep n+1
# (PUSH, idx of pending send)
idx = cast(int, task_path[1])
if idx >= len(checkpoint["pending_sends"]):
@@ -601,43 +543,6 @@ def prepare_single_task(
PUSH,
str(idx),
)
elif len(task_path) >= 4:
# new PUSH tasks, executed in superstep n
# (PUSH, parent task path, idx of PUSH write, id of parent task)
task_path_tt = cast(tuple[str, tuple, int, str], task_path)
writes_for_path = [w for w in pending_writes if w[0] == task_path_tt[3]]
if task_path_tt[2] >= len(writes_for_path):
logger.warning(
f"Ignoring invalid write index {task_path[2]} in pending writes"
)
return
packet = writes_for_path[task_path_tt[2]][2]
if packet is None:
return
if not isinstance(packet, Send):
logger.warning(
f"Ignoring invalid packet type {type(packet)} in pending writes"
)
return
if packet.node not in processes:
logger.warning(
f"Ignoring unknown node name {packet.node} in pending writes"
)
return
# create task id
triggers = [PUSH]
checkpoint_ns = (
f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node
)
task_id = _uuid5_str(
checkpoint_id,
checkpoint_ns,
str(step),
packet.node,
PUSH,
_tuple_str(task_path[1]),
str(task_path[2]),
)
else:
logger.warning(f"Ignoring invalid PUSH task path {task_path}")
return
@@ -904,7 +809,9 @@ def _uuid5_str(namespace: bytes, *parts: str) -> str:
def _tuple_str(tup: Union[str, int, tuple]) -> str:
"""Generate a string representation of a tuple."""
return (
f"({', '.join(_tuple_str(x) for x in tup)})"
f"~{', '.join(_tuple_str(x) for x in tup)}"
if isinstance(tup, (tuple, list))
else f"{tup:010d}"
if isinstance(tup, int)
else str(tup)
)
+1 -3
View File
@@ -8,10 +8,8 @@ from langgraph.checkpoint.base import PendingWrite
from langgraph.constants import (
EMPTY_SEQ,
ERROR,
FF_SEND_V2,
INTERRUPT,
NULL_TASK_ID,
PUSH,
RESUME,
RETURN,
SELF,
@@ -82,7 +80,7 @@ def map_command(
sends = [cmd.goto]
for send in sends:
if isinstance(send, Send):
yield (NULL_TASK_ID, PUSH if FF_SEND_V2 else TASKS, send)
yield (NULL_TASK_ID, TASKS, send)
elif isinstance(send, str):
yield (NULL_TASK_ID, f"branch:{START}:{SELF}:{send}", START)
else:
+5 -3
View File
@@ -693,9 +693,6 @@ class PregelLoop(LoopProtocol):
traceback: Optional[TracebackType],
) -> Optional[bool]:
suppress = isinstance(exc_value, GraphInterrupt) and not self.is_nested
if suppress or exc_type is None:
# save final output
self.output = read_channels(self.channels, self.output_keys)
if suppress:
# emit one last "values" event, with pending writes applied
if (
@@ -723,8 +720,13 @@ class PregelLoop(LoopProtocol):
"updates",
lambda: iter([{INTERRUPT: cast(GraphInterrupt, exc_value).args[0]}]),
)
# save final output
self.output = read_channels(self.channels, self.output_keys)
# suppress interrupt
return True
elif exc_type is None:
# save final output
self.output = read_channels(self.channels, self.output_keys)
def _emit(
self,
+3 -3
View File
@@ -14,7 +14,7 @@ from typing import (
from langchain_core.runnables import Runnable, RunnableConfig
from langchain_core.runnables.utils import ConfigurableFieldSpec
from langgraph.constants import CONF, CONFIG_KEY_SEND, FF_SEND_V2, PUSH, TASKS, Send
from langgraph.constants import CONF, CONFIG_KEY_SEND, TASKS, Send
from langgraph.errors import InvalidUpdateError
from langgraph.utils.runnable import RunnableCallable
@@ -125,7 +125,7 @@ class ChannelWrite(RunnableCallable):
# validate
for w in writes:
if isinstance(w, ChannelWriteEntry):
if w.channel in (TASKS, PUSH):
if w.channel == TASKS:
raise InvalidUpdateError(
"Cannot write to the reserved channel TASKS"
)
@@ -138,7 +138,7 @@ class ChannelWrite(RunnableCallable):
tuples: list[tuple[str, Any]] = []
for w in writes:
if isinstance(w, Send):
tuples.append((PUSH if FF_SEND_V2 else TASKS, w))
tuples.append((TASKS, w))
elif isinstance(w, ChannelWriteTupleEntry):
if ww := w.mapper(w.value):
tuples.extend(ww)
+25 -1
View File
@@ -1,5 +1,6 @@
from langgraph.checkpoint.base import empty_checkpoint
from langgraph.pregel.algo import prepare_next_tasks
from langgraph.constants import PULL, PUSH
from langgraph.pregel.algo import _tuple_str, prepare_next_tasks
from langgraph.pregel.manager import ChannelsManager
@@ -40,3 +41,26 @@ def test_prepare_next_tasks() -> None:
)
# TODO: add more tests
def test_tuple_str() -> None:
push_path_a = (PUSH, 2)
pull_path_a = (PULL, "abc")
push_path_b = (PUSH, push_path_a, 1)
push_path_c = (PUSH, push_path_b, 3)
assert _tuple_str(push_path_a) == f"~{PUSH}, 0000000002"
assert _tuple_str(push_path_b) == f"~{PUSH}, ~{PUSH}, 0000000002, 0000000001"
assert (
_tuple_str(push_path_c)
== f"~{PUSH}, ~{PUSH}, ~{PUSH}, 0000000002, 0000000001, 0000000003"
)
assert _tuple_str(pull_path_a) == f"~{PULL}, abc"
path_list = [push_path_b, push_path_a, pull_path_a, push_path_c]
assert sorted(map(_tuple_str, path_list)) == [
f"~{PULL}, abc",
f"~{PUSH}, 0000000002",
f"~{PUSH}, ~{PUSH}, 0000000002, 0000000001",
f"~{PUSH}, ~{PUSH}, ~{PUSH}, 0000000002, 0000000001, 0000000003",
]
+175 -553
View File
@@ -16,7 +16,7 @@ from langgraph.channels.context import Context
from langgraph.channels.last_value import LastValue
from langgraph.channels.untracked_value import UntrackedValue
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.constants import END, FF_SEND_V2, PULL, PUSH, START
from langgraph.constants import END, PULL, PUSH, START
from langgraph.errors import NodeInterrupt
from langgraph.graph import StateGraph
from langgraph.graph.graph import Graph
@@ -3021,9 +3021,6 @@ def test_state_graph_packets(
{"__interrupt__": ()},
]
if not FF_SEND_V2:
return
assert app_w_interrupt.get_state(config) == StateSnapshot(
values={
"messages": [
@@ -3041,33 +3038,7 @@ def test_state_graph_packets(
),
]
},
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
"messages": AIMessage(
content="",
additional_kwargs={},
response_metadata={},
id="ai1",
tool_calls=[
{
"name": "search_api",
"args": {"query": "query"},
"id": "tool_call123",
"type": "tool_call",
}
],
)
},
),
PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)),
),
tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),),
next=("tools",),
config={
"configurable": {
@@ -3080,8 +3051,23 @@ def test_state_graph_packets(
metadata={
"parents": {},
"source": "loop",
"step": 0,
"writes": None,
"step": 1,
"writes": {
"agent": {
"messages": AIMessage(
content="",
id="ai1",
tool_calls=[
{
"name": "search_api",
"args": {"query": "query"},
"id": "tool_call123",
"type": "tool_call",
}
],
)
}
},
"thread_id": "1",
},
parent_config=(
@@ -3116,7 +3102,7 @@ def test_state_graph_packets(
),
]
},
tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0)),),
tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),),
next=("tools",),
config={
"configurable": {
@@ -3129,7 +3115,7 @@ def test_state_graph_packets(
metadata={
"parents": {},
"source": "update",
"step": 1,
"step": 2,
"writes": {
"agent": {
"messages": AIMessage(
@@ -3227,16 +3213,26 @@ def test_state_graph_packets(
]
},
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
PregelTask(AnyStr(), "tools", (PUSH, 0)),
PregelTask(AnyStr(), "tools", (PUSH, 1)),
),
next=("tools", "tools"),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
created_at=AnyStr(),
metadata={
"parents": {},
"source": "loop",
"step": 4,
"writes": {
"agent": {
"messages": AIMessage(
"",
content="",
id="ai2",
tool_calls=[
{
@@ -3254,31 +3250,6 @@ def test_state_graph_packets(
],
)
},
),
PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)),
PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)),
),
next=("tools", "tools"),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
created_at=AnyStr(),
metadata={
"parents": {},
"source": "loop",
"step": 2,
"writes": {
"tools": {
"messages": _AnyIdToolMessage(
content="result for a different query",
name="search_api",
tool_call_id="tool_call123",
),
},
},
"thread_id": "1",
},
@@ -3334,7 +3305,7 @@ def test_state_graph_packets(
metadata={
"parents": {},
"source": "update",
"step": 3,
"step": 5,
"writes": {
"agent": {
"messages": AIMessage(content="answer", id="ai2"),
@@ -3400,31 +3371,7 @@ def test_state_graph_packets(
),
]
},
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
"messages": AIMessage(
"",
id="ai1",
tool_calls=[
{
"name": "search_api",
"args": {"query": "query"},
"id": "tool_call123",
"type": "tool_call",
}
],
)
},
),
PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)),
),
tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),),
next=("tools",),
config={
"configurable": {
@@ -3437,8 +3384,23 @@ def test_state_graph_packets(
metadata={
"parents": {},
"source": "loop",
"step": 0,
"writes": None,
"step": 1,
"writes": {
"agent": {
"messages": AIMessage(
content="",
id="ai1",
tool_calls=[
{
"name": "search_api",
"args": {"query": "query"},
"id": "tool_call123",
"type": "tool_call",
}
],
)
}
},
"thread_id": "2",
},
parent_config=(
@@ -3473,14 +3435,14 @@ def test_state_graph_packets(
),
]
},
tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0)),),
tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=AnyStr(),
metadata={
"parents": {},
"source": "update",
"step": 1,
"step": 2,
"writes": {
"agent": {
"messages": AIMessage(
@@ -3578,36 +3540,8 @@ def test_state_graph_packets(
]
},
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
"messages": AIMessage(
"",
id="ai2",
tool_calls=[
{
"name": "search_api",
"args": {"query": "another", "idx": 0},
"id": "tool_call234",
"type": "tool_call",
},
{
"name": "search_api",
"args": {"query": "a third one", "idx": 1},
"id": "tool_call567",
"type": "tool_call",
},
],
)
},
),
PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)),
PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)),
PregelTask(AnyStr(), "tools", (PUSH, 0)),
PregelTask(AnyStr(), "tools", (PUSH, 1)),
),
next=("tools", "tools"),
config={
@@ -3621,14 +3555,25 @@ def test_state_graph_packets(
metadata={
"parents": {},
"source": "loop",
"step": 2,
"step": 4,
"writes": {
"tools": {
"messages": _AnyIdToolMessage(
content="result for a different query",
name="search_api",
tool_call_id="tool_call123",
),
"agent": {
"messages": AIMessage(
id="ai2",
content="",
tool_calls=[
{
"id": "tool_call234",
"name": "search_api",
"args": {"query": "another", "idx": 0},
},
{
"id": "tool_call567",
"name": "search_api",
"args": {"query": "a third one", "idx": 1},
},
],
)
},
},
"thread_id": "2",
@@ -3685,7 +3630,7 @@ def test_state_graph_packets(
metadata={
"parents": {},
"source": "update",
"step": 3,
"step": 5,
"writes": {
"agent": {
"messages": AIMessage(content="answer", id="ai2"),
@@ -5859,7 +5804,6 @@ def test_dynamic_interrupt(
)
@pytest.mark.skipif(not FF_SEND_V2, reason="send v2 is not enabled")
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_copy_checkpoint(
request: pytest.FixtureRequest, checkpointer_name: str
@@ -5879,6 +5823,7 @@ def test_copy_checkpoint(
nonlocal tool_two_node_count
tool_two_node_count += 1
if s["market"] == "DE":
time.sleep(0.1)
answer = interrupt("Just because...")
else:
answer = " all good"
@@ -5908,7 +5853,7 @@ def test_copy_checkpoint(
assert run.outputs == {"market": "DE", "my_key": "value one"}
assert tool_two.invoke({"my_key": "value", "market": "US"}) == {
"my_key": "value one all good",
"my_key": "value all good one",
"market": "US",
}
@@ -5939,6 +5884,7 @@ def test_copy_checkpoint(
]
# resume with answer
assert [c for c in tool_two.stream(Command(resume=" my answer"), thread2)] == [
{"tool_one": {"my_key": " one"}, "__metadata__": {"cached": True}},
{"tool_two": {"my_key": " my answer"}},
]
@@ -5955,7 +5901,7 @@ def test_copy_checkpoint(
"parents": {},
"source": "loop",
"step": 0,
"writes": {"tool_one": {"my_key": " one"}},
"writes": None,
"thread_id": "1",
},
{
@@ -5971,6 +5917,12 @@ def test_copy_checkpoint(
values={"my_key": "value ⛰️ one", "market": "DE"},
next=("tool_two",),
tasks=(
PregelTask(
id=AnyStr(),
name="tool_one",
path=("__pregel_push", 0),
result={"my_key": " one"},
),
PregelTask(
AnyStr(),
"tool_two",
@@ -5996,7 +5948,7 @@ def test_copy_checkpoint(
"parents": {},
"source": "loop",
"step": 0,
"writes": {"tool_one": {"my_key": " one"}},
"writes": None,
"thread_id": "1",
},
parent_config=(
@@ -6005,13 +5957,25 @@ def test_copy_checkpoint(
else [*tool_two.checkpointer.list(thread1, limit=2)][-1].config
),
)
if "shallow" in checkpointer_name:
return
# clear the interrupt and next tasks
tool_two.update_state(thread1, None)
tool_two.update_state(thread1, None, as_node="__copy__")
# interrupt is cleared, next task is kept
assert tool_two.get_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️ one", "market": "DE"},
next=("tool_two",),
values={"my_key": "value ⛰️", "market": "DE"},
next=(
"tool_one",
"tool_two",
),
tasks=(
PregelTask(
id=AnyStr(),
name="tool_one",
path=("__pregel_push", 0),
),
PregelTask(
AnyStr(),
"tool_two",
@@ -6029,15 +5993,13 @@ def test_copy_checkpoint(
created_at=AnyStr(),
metadata={
"parents": {},
"source": "update",
"source": "fork",
"step": 1,
"writes": {},
"writes": None,
"thread_id": "1",
},
parent_config=(
None
if "shallow" in checkpointer_name
else [*tool_two.checkpointer.list(thread1, limit=2)][-1].config
[*tool_two.checkpointer.list(thread1, limit=2)][-1].parent_config
),
)
@@ -7292,12 +7254,11 @@ def test_branch_then(
)
@pytest.mark.skip("TODO: re-enable in next PR")
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_send_dedupe_on_resume(
request: pytest.FixtureRequest, checkpointer_name: str
) -> None:
if not FF_SEND_V2:
pytest.skip("Send deduplication is only available in Send V2")
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
class InterruptOnce:
@@ -9320,328 +9281,18 @@ def test_send_to_nested_graphs(
# check state
outer_state = graph.get_state(config)
if not FF_SEND_V2:
# update state of dogs joke graph
graph.update_state(outer_state.tasks[1].state, {"subject": "turtles - hohoho"})
# continue past interrupt
assert sorted(
graph.stream(None, config=config),
key=lambda d: d["generate_joke"]["jokes"][0],
) == [
{"generate_joke": {"jokes": ["Joke about cats - hohoho"]}},
{"generate_joke": {"jokes": ["Joke about turtles - hohoho"]}},
]
return
assert outer_state == StateSnapshot(
values={"subjects": ["cats", "dogs"], "jokes": []},
tasks=(
PregelTask(
id=AnyStr(),
name="__start__",
path=("__pregel_pull", "__start__"),
error=None,
interrupts=(),
state=None,
result={"subjects": ["cats", "dogs"]},
),
PregelTask(
AnyStr(),
"generate_joke",
(PUSH, ("__pregel_pull", "__start__"), 1),
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
),
PregelTask(
AnyStr(),
"generate_joke",
(PUSH, ("__pregel_pull", "__start__"), 2),
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
),
),
next=("generate_joke", "generate_joke"),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "input",
"writes": {"__start__": {"subjects": ["cats", "dogs"]}},
"step": -1,
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=None,
)
# check state of each of the inner tasks
assert graph.get_state(outer_state.tasks[1].state) == StateSnapshot(
values={"subject": "cats - hohoho", "jokes": []},
next=("generate",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("generate_joke:"): AnyStr(),
}
),
}
},
metadata={
"step": 1,
"source": "loop",
"writes": None,
"parents": {"": AnyStr()},
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
"langgraph_checkpoint_ns": AnyStr("generate_joke:"),
"langgraph_node": "generate_joke",
"langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 1],
"langgraph_step": 0,
"langgraph_triggers": [PUSH],
},
created_at=AnyStr(),
parent_config=(
None
if "shallow" in checkpointer_name
else {
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("generate_joke:"): AnyStr(),
}
),
}
}
),
tasks=(PregelTask(id=AnyStr(""), name="generate", path=(PULL, "generate")),),
)
assert graph.get_state(outer_state.tasks[2].state) == StateSnapshot(
values={"subject": "dogs - hohoho", "jokes": []},
next=("generate",),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("generate_joke:"): AnyStr(),
}
),
}
},
metadata={
"step": 1,
"source": "loop",
"writes": None,
"parents": {"": AnyStr()},
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
"langgraph_checkpoint_ns": AnyStr("generate_joke:"),
"langgraph_node": "generate_joke",
"langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 2],
"langgraph_step": 0,
"langgraph_triggers": [PUSH],
},
created_at=AnyStr(),
parent_config=(
None
if "shallow" in checkpointer_name
else {
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
"checkpoint_id": AnyStr(),
"checkpoint_map": AnyDict(
{
"": AnyStr(),
AnyStr("generate_joke:"): AnyStr(),
}
),
}
}
),
tasks=(PregelTask(id=AnyStr(""), name="generate", path=(PULL, "generate")),),
)
# update state of dogs joke graph
graph.update_state(
outer_state.tasks[2 if FF_SEND_V2 else 1].state, {"subject": "turtles - hohoho"}
)
graph.update_state(outer_state.tasks[1].state, {"subject": "turtles - hohoho"})
# continue past interrupt
assert sorted(
graph.stream(None, config=config), key=lambda d: d["generate_joke"]["jokes"][0]
graph.stream(None, config=config),
key=lambda d: d["generate_joke"]["jokes"][0],
) == [
{"generate_joke": {"jokes": ["Joke about cats - hohoho"]}},
{"generate_joke": {"jokes": ["Joke about turtles - hohoho"]}},
]
actual_snapshot = graph.get_state(config)
expected_snapshot = StateSnapshot(
values={
"subjects": ["cats", "dogs"],
"jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"],
},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"writes": {
"generate_joke": [
{"jokes": ["Joke about cats - hohoho"]},
{"jokes": ["Joke about turtles - hohoho"]},
]
},
"step": 0,
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=(
None
if "shallow" in checkpointer_name
else {
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
}
),
)
assert actual_snapshot == expected_snapshot
if "shallow" in checkpointer_name:
return
# test full history
actual_history = list(graph.get_state_history(config))
# get subgraph node state for expected history
expected_history = [
StateSnapshot(
values={
"subjects": ["cats", "dogs"],
"jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"],
},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"writes": {
"generate_joke": [
{"jokes": ["Joke about cats - hohoho"]},
{"jokes": ["Joke about turtles - hohoho"]},
]
},
"step": 0,
"thread_id": "1",
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
),
StateSnapshot(
values={"jokes": []},
tasks=(
PregelTask(
id=AnyStr(),
name="__start__",
path=("__pregel_pull", "__start__"),
error=None,
interrupts=(),
state=None,
result={"subjects": ["cats", "dogs"]},
),
PregelTask(
AnyStr(),
"generate_joke",
(PUSH, ("__pregel_pull", "__start__"), 1),
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
result={"jokes": ["Joke about cats - hohoho"]},
),
PregelTask(
AnyStr(),
"generate_joke",
(PUSH, ("__pregel_pull", "__start__"), 2),
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
result={"jokes": ["Joke about turtles - hohoho"]},
),
),
next=("__start__", "generate_joke", "generate_joke"),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "input",
"writes": {"__start__": {"subjects": ["cats", "dogs"]}},
"step": -1,
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=None,
),
]
assert actual_history == expected_history
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_send_react_interrupt(
@@ -9767,9 +9418,6 @@ def test_send_react_interrupt(
}
assert foo_called == 0
if not FF_SEND_V2:
return
# get state should show the pending task
state = graph.get_state(thread1)
assert state == StateSnapshot(
@@ -9798,9 +9446,24 @@ def test_send_react_interrupt(
}
},
metadata={
"step": 0,
"step": 1,
"source": "loop",
"writes": None,
"writes": {
"agent": {
"messages": AIMessage(
content="",
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
)
}
},
"parents": {},
"thread_id": "2",
},
@@ -9817,34 +9480,10 @@ def test_send_react_interrupt(
}
),
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
"messages": AIMessage(
content="",
additional_kwargs={},
response_metadata={},
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
)
},
),
PregelTask(
id=AnyStr(),
name="foo",
path=("__pregel_push", ("__pregel_pull", "agent"), 2),
path=("__pregel_push", 0),
error=None,
interrupts=(),
state=None,
@@ -9878,7 +9517,7 @@ def test_send_react_interrupt(
}
},
metadata={
"step": 1,
"step": 2,
"source": "update",
"writes": {
"agent": {
@@ -9965,9 +9604,26 @@ def test_send_react_interrupt(
}
},
metadata={
"step": 0,
"step": 1,
"source": "loop",
"writes": None,
"writes": {
"agent": {
"messages": AIMessage(
content="",
additional_kwargs={},
response_metadata={},
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
)
}
},
"parents": {},
"thread_id": "3",
},
@@ -9984,32 +9640,10 @@ def test_send_react_interrupt(
}
),
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
"messages": AIMessage(
"",
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
)
},
),
PregelTask(
id=AnyStr(),
name="foo",
path=("__pregel_push", ("__pregel_pull", "agent"), 2),
path=("__pregel_push", 0),
error=None,
interrupts=(),
state=None,
@@ -10064,7 +9698,7 @@ def test_send_react_interrupt(
}
},
metadata={
"step": 1,
"step": 2,
"source": "update",
"writes": {
"agent": {
@@ -10100,7 +9734,7 @@ def test_send_react_interrupt(
PregelTask(
id=AnyStr(),
name="foo",
path=("__pregel_push", (), 0),
path=("__pregel_push", 0),
error=None,
interrupts=(),
state=None,
@@ -10230,9 +9864,6 @@ def test_send_react_interrupt_control(
}
assert foo_called == 1
if not FF_SEND_V2:
return
# interrupt-update-resume flow
foo_called = 0
graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"])
@@ -10283,9 +9914,24 @@ def test_send_react_interrupt_control(
}
},
metadata={
"step": 0,
"step": 1,
"source": "loop",
"writes": None,
"writes": {
"agent": {
"messages": AIMessage(
content="",
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
)
}
},
"parents": {},
"thread_id": "2",
},
@@ -10302,34 +9948,10 @@ def test_send_react_interrupt_control(
}
),
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
"messages": AIMessage(
content="",
additional_kwargs={},
response_metadata={},
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
)
},
),
PregelTask(
id=AnyStr(),
name="foo",
path=("__pregel_push", ("__pregel_pull", "agent"), 2),
path=("__pregel_push", 0),
error=None,
interrupts=(),
state=None,
@@ -10363,7 +9985,7 @@ def test_send_react_interrupt_control(
}
},
metadata={
"step": 1,
"step": 2,
"source": "update",
"writes": {
"agent": {
+76 -351
View File
@@ -25,7 +25,7 @@ from syrupy import SnapshotAssertion
from langgraph.channels.context import Context
from langgraph.channels.last_value import LastValue
from langgraph.channels.untracked_value import UntrackedValue
from langgraph.constants import END, FF_SEND_V2, PULL, PUSH, START
from langgraph.constants import END, PULL, PUSH, START
from langgraph.graph.graph import Graph
from langgraph.graph.message import MessageGraph, add_messages
from langgraph.graph.state import StateGraph
@@ -2740,9 +2740,6 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
{"__interrupt__": ()},
]
if not FF_SEND_V2:
return
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
values={
"messages": [
@@ -2760,31 +2757,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
),
]
},
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
"messages": AIMessage(
"",
id="ai1",
tool_calls=[
{
"name": "search_api",
"args": {"query": "query"},
"id": "tool_call123",
"type": "tool_call",
}
],
)
},
),
PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)),
),
tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),),
next=("tools",),
config={
"configurable": {
@@ -2797,8 +2770,23 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
metadata={
"parents": {},
"source": "loop",
"step": 0,
"writes": None,
"step": 1,
"writes": {
"agent": {
"messages": AIMessage(
content="",
id="ai1",
tool_calls=[
{
"name": "search_api",
"args": {"query": "query"},
"id": "tool_call123",
"type": "tool_call",
}
],
)
}
},
"thread_id": "1",
},
parent_config=(
@@ -2834,14 +2822,14 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
),
]
},
tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0)),),
tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),),
next=("tools",),
config=tup.config,
created_at=tup.checkpoint["ts"],
metadata={
"parents": {},
"source": "update",
"step": 1,
"step": 2,
"writes": {
"agent": {
"messages": AIMessage(
@@ -2941,36 +2929,8 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
]
},
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
"messages": AIMessage(
"",
id="ai2",
tool_calls=[
{
"name": "search_api",
"args": {"query": "another", "idx": 0},
"id": "tool_call234",
"type": "tool_call",
},
{
"name": "search_api",
"args": {"query": "a third one", "idx": 1},
"id": "tool_call567",
"type": "tool_call",
},
],
)
},
),
PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)),
PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)),
PregelTask(AnyStr(), "tools", (PUSH, 0)),
PregelTask(AnyStr(), "tools", (PUSH, 1)),
),
next=("tools", "tools"),
config=tup.config,
@@ -2978,13 +2938,24 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
metadata={
"parents": {},
"source": "loop",
"step": 2,
"step": 4,
"writes": {
"tools": {
"messages": _AnyIdToolMessage(
content="result for a different query",
name="search_api",
tool_call_id="tool_call123",
"agent": {
"messages": AIMessage(
id="ai2",
content="",
tool_calls=[
{
"id": "tool_call234",
"name": "search_api",
"args": {"query": "another", "idx": 0},
},
{
"id": "tool_call567",
"name": "search_api",
"args": {"query": "a third one", "idx": 1},
},
],
),
},
},
@@ -3036,7 +3007,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
metadata={
"parents": {},
"source": "update",
"step": 3,
"step": 5,
"writes": {
"agent": {
"messages": AIMessage(content="answer", id="ai2"),
@@ -3103,15 +3074,16 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
),
]
},
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),),
next=("tools",),
config=tup.config,
created_at=tup.checkpoint["ts"],
metadata={
"parents": {},
"source": "loop",
"step": 1,
"writes": {
"agent": {
"messages": AIMessage(
content="",
additional_kwargs={},
@@ -3126,18 +3098,8 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
}
],
)
},
),
PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)),
),
next=("tools",),
config=tup.config,
created_at=tup.checkpoint["ts"],
metadata={
"parents": {},
"source": "loop",
"step": 0,
"writes": None,
}
},
"thread_id": "2",
},
parent_config=(
@@ -3173,14 +3135,14 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
),
]
},
tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0)),),
tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),),
next=("tools",),
config=tup.config,
created_at=tup.checkpoint["ts"],
metadata={
"parents": {},
"source": "update",
"step": 1,
"step": 2,
"writes": {
"agent": {
"messages": AIMessage(
@@ -3280,38 +3242,8 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
]
},
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
"messages": AIMessage(
content="",
additional_kwargs={},
response_metadata={},
id="ai2",
tool_calls=[
{
"name": "search_api",
"args": {"query": "another", "idx": 0},
"id": "tool_call234",
"type": "tool_call",
},
{
"name": "search_api",
"args": {"query": "a third one", "idx": 1},
"id": "tool_call567",
"type": "tool_call",
},
],
)
},
),
PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)),
PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)),
PregelTask(AnyStr(), "tools", (PUSH, 0)),
PregelTask(AnyStr(), "tools", (PUSH, 1)),
),
next=("tools", "tools"),
config=tup.config,
@@ -3319,13 +3251,24 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
metadata={
"parents": {},
"source": "loop",
"step": 2,
"step": 4,
"writes": {
"tools": {
"messages": _AnyIdToolMessage(
content="result for a different query",
name="search_api",
tool_call_id="tool_call123",
"agent": {
"messages": AIMessage(
id="ai2",
content="",
tool_calls=[
{
"id": "tool_call234",
"name": "search_api",
"args": {"query": "another", "idx": 0},
},
{
"id": "tool_call567",
"name": "search_api",
"args": {"query": "a third one", "idx": 1},
},
],
),
},
},
@@ -3377,7 +3320,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None:
metadata={
"parents": {},
"source": "update",
"step": 3,
"step": 5,
"writes": {
"agent": {
"messages": AIMessage(content="answer", id="ai2"),
@@ -6961,83 +6904,9 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None:
# check state
outer_state = await graph.aget_state(config)
if not FF_SEND_V2:
# update state of dogs joke graph
await graph.aupdate_state(
outer_state.tasks[1].state, {"subject": "turtles - hohoho"}
)
# continue past interrupt
assert await graph.ainvoke(None, config=config) == {
"subjects": ["cats", "dogs"],
"jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"],
}
return
assert outer_state == StateSnapshot(
values={"subjects": ["cats", "dogs"], "jokes": []},
tasks=(
PregelTask(
id=AnyStr(),
name="__start__",
path=("__pregel_pull", "__start__"),
error=None,
interrupts=(),
state=None,
result={"subjects": ["cats", "dogs"]},
),
PregelTask(
AnyStr(),
"generate_joke",
(PUSH, ("__pregel_pull", "__start__"), 1),
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
),
PregelTask(
AnyStr(),
"generate_joke",
(PUSH, ("__pregel_pull", "__start__"), 2),
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
),
),
next=("generate_joke", "generate_joke"),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "input",
"writes": {
"__start__": {
"subjects": [
"cats",
"dogs",
],
}
},
"step": -1,
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=None,
)
# update state of dogs joke graph
await graph.aupdate_state(
outer_state.tasks[2].state, {"subject": "turtles - hohoho"}
outer_state.tasks[1].state, {"subject": "turtles - hohoho"}
)
# continue past interrupt
@@ -7046,150 +6915,6 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None:
"jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"],
}
actual_snapshot = await graph.aget_state(config)
expected_snapshot = StateSnapshot(
values={
"subjects": ["cats", "dogs"],
"jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"],
},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"writes": {
"generate_joke": [
{"jokes": ["Joke about cats - hohoho"]},
{"jokes": ["Joke about turtles - hohoho"]},
]
},
"step": 0,
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=(
None
if "shallow" in checkpointer_name
else {
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
}
),
)
assert actual_snapshot == expected_snapshot
if "shallow" in checkpointer_name:
return
# test full history
actual_history = [c async for c in graph.aget_state_history(config)]
expected_history = [
StateSnapshot(
values={
"subjects": ["cats", "dogs"],
"jokes": [
"Joke about cats - hohoho",
"Joke about turtles - hohoho",
],
},
tasks=(),
next=(),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "loop",
"writes": {
"generate_joke": [
{"jokes": ["Joke about cats - hohoho"]},
{"jokes": ["Joke about turtles - hohoho"]},
]
},
"step": 0,
"thread_id": "1",
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
),
StateSnapshot(
values={"jokes": []},
next=("__start__", "generate_joke", "generate_joke"),
tasks=(
PregelTask(
id=AnyStr(),
name="__start__",
path=("__pregel_pull", "__start__"),
error=None,
interrupts=(),
state=None,
result={"subjects": ["cats", "dogs"]},
),
PregelTask(
AnyStr(),
"generate_joke",
(PUSH, ("__pregel_pull", "__start__"), 1),
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
result={"jokes": ["Joke about cats - hohoho"]},
),
PregelTask(
AnyStr(),
"generate_joke",
(PUSH, ("__pregel_pull", "__start__"), 2),
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
result={"jokes": ["Joke about turtles - hohoho"]},
),
),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"parents": {},
"source": "input",
"writes": {"__start__": {"subjects": ["cats", "dogs"]}},
"step": -1,
"thread_id": "1",
},
created_at=AnyStr(),
parent_config=None,
),
]
assert actual_history == expected_history
@pytest.mark.skipif(
sys.version_info < (3, 11),
+23 -57
View File
@@ -50,13 +50,7 @@ from langgraph.checkpoint.base import (
CheckpointTuple,
)
from langgraph.checkpoint.memory import MemorySaver
from langgraph.constants import (
CONFIG_KEY_NODE_FINISHED,
ERROR,
FF_SEND_V2,
PULL,
START,
)
from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, START
from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError
from langgraph.func import entrypoint, task
from langgraph.graph import END, Graph, StateGraph
@@ -1375,31 +1369,17 @@ def test_concurrent_emit_sends() -> None:
builder.add_conditional_edges("1.1", send_for_profit)
builder.add_conditional_edges("2", route_to_three)
graph = builder.compile()
assert graph.invoke(["0"]) == (
[
"0",
"1",
"1.1",
"2|1",
"2|2",
"2|3",
"2|4",
"3",
"3.1",
]
if FF_SEND_V2
else [
"0",
"1",
"1.1",
"3.1",
"2|1",
"2|2",
"2|3",
"2|4",
"3",
]
)
assert graph.invoke(["0"]) == [
"0",
"1",
"1.1",
"3.1",
"2|1",
"2|2",
"2|3",
"2|4",
"3",
]
def test_send_sequences() -> None:
@@ -1438,31 +1418,17 @@ def test_send_sequences() -> None:
builder.add_conditional_edges("1", send_for_fun)
builder.add_conditional_edges("2", route_to_three)
graph = builder.compile()
assert (
graph.invoke(["0"])
== [
"0",
"1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='2', arg=4))",
"2|3",
"2|4",
"3",
"3.1",
]
if FF_SEND_V2
else [
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='2', arg=4))",
"3",
"2|3",
"2|4",
"3",
]
)
assert graph.invoke(["0"]) == [
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='2', arg=4))",
"3",
"2|3",
"2|4",
"3",
]
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
+250 -226
View File
@@ -47,13 +47,7 @@ from langgraph.checkpoint.base import (
CheckpointTuple,
)
from langgraph.checkpoint.memory import MemorySaver
from langgraph.constants import (
CONFIG_KEY_NODE_FINISHED,
ERROR,
FF_SEND_V2,
PULL,
START,
)
from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH, START
from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt
from langgraph.func import entrypoint, task
from langgraph.graph import END, Graph, StateGraph
@@ -878,7 +872,6 @@ async def test_dynamic_interrupt_subgraph(checkpointer_name: str) -> None:
)
@pytest.mark.skipif(not FF_SEND_V2, reason="send v2 is not enabled")
@pytest.mark.skipif(
sys.version_info < (3, 11),
reason="Python 3.11+ is required for async contextvars support",
@@ -914,7 +907,7 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
tracer = FakeTracer()
assert await tool_two.ainvoke(
{"my_key": "value", "market": "DE"}, {"callbacks": [tracer]}
{"my_key": "value", "market": "DE"}, {"callbacks": [tracer]}, debug=True
) == {
"my_key": "value one",
"market": "DE",
@@ -927,7 +920,7 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
assert run.outputs == {"market": "DE", "my_key": "value one"}
assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == {
"my_key": "value one all good",
"my_key": "value all good one",
"market": "US",
}
@@ -964,6 +957,10 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
assert [
c async for c in tool_two.astream(Command(resume=" my answer"), thread2)
] == [
{
"__metadata__": {"cached": True},
"tool_one": {"my_key": " one"},
},
{"tool_two": {"my_key": " my answer"}},
]
@@ -983,7 +980,7 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
"parents": {},
"source": "loop",
"step": 0,
"writes": {"tool_one": {"my_key": " one"}},
"writes": None,
"thread_id": "1",
},
{
@@ -1000,6 +997,15 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
values={"my_key": "value ⛰️ one", "market": "DE"},
next=("tool_two",),
tasks=(
PregelTask(
AnyStr(),
name="tool_one",
path=("__pregel_push", 0),
error=None,
interrupts=(),
state=None,
result={"my_key": " one"},
),
PregelTask(
AnyStr(),
"tool_two",
@@ -1019,7 +1025,7 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
"parents": {},
"source": "loop",
"step": 0,
"writes": {"tool_one": {"my_key": " one"}},
"writes": None,
"thread_id": "1",
},
parent_config=(
@@ -1030,14 +1036,25 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
].config
),
)
if "shallow" in checkpointer_name:
# shallow checkpointer doesn't support copy
return
# clear the interrupt and next tasks
await tool_two.aupdate_state(thread1, None)
await tool_two.aupdate_state(thread1, None, as_node="__copy__")
# interrupt is cleared, next task is kept
tup = await tool_two.checkpointer.aget_tuple(thread1)
assert await tool_two.aget_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️ one", "market": "DE"},
next=("tool_two",),
values={"my_key": "value ⛰️", "market": "DE"},
next=("tool_one", "tool_two"),
tasks=(
PregelTask(
AnyStr(),
"tool_one",
(PUSH, 0),
result=None,
),
PregelTask(
AnyStr(),
"tool_two",
@@ -1049,17 +1066,15 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None:
created_at=tup.checkpoint["ts"],
metadata={
"parents": {},
"source": "update",
"source": "fork",
"step": 1,
"writes": {},
"writes": None,
"thread_id": "1",
},
parent_config=(
None
if "shallow" in checkpointer_name
else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][
[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][
-1
].config
].parent_config
),
)
@@ -2345,18 +2360,6 @@ async def test_concurrent_emit_sends() -> None:
graph = builder.compile()
assert await graph.ainvoke(["0"]) == (
[
"0",
"1",
"1.1",
"2|1",
"2|2",
"2|3",
"2|4",
"3",
"3.1",
]
if FF_SEND_V2
else [
"0",
"1",
"1.1",
@@ -2370,6 +2373,7 @@ async def test_concurrent_emit_sends() -> None:
)
@pytest.mark.skip("TODO: re-enable in next PR")
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_send_sequences(checkpointer_name: str) -> None:
class Node:
@@ -2407,34 +2411,17 @@ async def test_send_sequences(checkpointer_name: str) -> None:
builder.add_conditional_edges("1", send_for_fun)
builder.add_conditional_edges("2", route_to_three)
graph = builder.compile()
assert (
await graph.ainvoke(["0"])
== [
"0",
"1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='2', arg=4))",
"2|3",
"2|4",
"3",
"3.1",
]
if FF_SEND_V2
else [
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='2', arg=4))",
"3",
"2|3",
"2|4",
"3",
]
)
if not FF_SEND_V2:
return
assert await graph.ainvoke(["0"]) == [
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='2', arg=4))",
"3",
"2|3",
"2|4",
"3",
]
async with awith_checkpointer(checkpointer_name) as checkpointer:
graph = builder.compile(checkpointer=checkpointer, interrupt_before=["3.1"])
@@ -2442,20 +2429,17 @@ async def test_send_sequences(checkpointer_name: str) -> None:
assert await graph.ainvoke(["0"], thread1) == [
"0",
"1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='2', arg=4))",
"2|3",
"2|4",
]
assert await graph.ainvoke(None, thread1) == [
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='2', arg=4))",
"3",
"2|3",
"2|4",
"3",
"3.1",
]
@@ -2631,11 +2615,9 @@ async def test_imp_stream_order(checkpointer_name: str) -> None:
]
@pytest.mark.skip("TODO: re-enable in next PR")
@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC)
async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
if not FF_SEND_V2:
pytest.skip("Send deduplication is only available in Send V2")
class InterruptOnce:
ticks: int = 0
@@ -2689,22 +2671,26 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
assert await graph.ainvoke(["0"], thread1, debug=1) == [
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"3",
"2|3",
]
assert builder.nodes["2"].runnable.func.ticks == 3
assert builder.nodes["flaky"].runnable.func.ticks == 1
print((await graph.aget_state(thread1)).tasks)
# resume execution
assert await graph.ainvoke(None, thread1, debug=1) == [
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"2|3",
"flaky|4",
"3",
"3.1",
"flaky|4",
"2|3",
"3",
]
# node "2" doesn't get called again, as we recover writes saved before
assert builder.nodes["2"].runnable.func.ticks == 3
@@ -2717,12 +2703,13 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
values=[
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"2|3",
"flaky|4",
"3",
"3.1",
"flaky|4",
"2|3",
"3",
],
next=(),
config={
@@ -2734,9 +2721,9 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
},
metadata={
"source": "loop",
"writes": {"3": ["3"], "3.1": ["3.1"]},
"writes": {"3": ["3"]},
"thread_id": "1",
"step": 2,
"step": 4,
"parents": {},
},
created_at=AnyStr(),
@@ -2753,12 +2740,14 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
values=[
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
"2|3",
"3",
"flaky|4",
"2|3",
],
next=("3", "3.1"),
next=("3",),
config={
"configurable": {
"thread_id": "1",
@@ -2768,17 +2757,9 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
},
metadata={
"source": "loop",
"writes": {
"1": ["1"],
"2": [
["2|Command(goto=Send(node='2', arg=3))"],
["2|Command(goto=Send(node='flaky', arg=4))"],
["2|3"],
],
"flaky": ["flaky|4"],
},
"writes": {"2": ["2|3"], "3": ["3"], "flaky": ["flaky|4"]},
"thread_id": "1",
"step": 1,
"step": 3,
"parents": {},
},
created_at=AnyStr(),
@@ -2799,6 +2780,123 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
state=None,
result=["3"],
),
),
),
StateSnapshot(
values=[
"0",
"1",
"3.1",
"2|Command(goto=Send(node='2', arg=3))",
"2|Command(goto=Send(node='flaky', arg=4))",
],
next=("2", "flaky", "3"),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {
"2": [
["2|Command(goto=Send(node='2', arg=3))"],
["2|Command(goto=Send(node='flaky', arg=4))"],
],
"3.1": ["3.1"],
},
"thread_id": "1",
"step": 2,
"parents": {},
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
tasks=(
PregelTask(
id=AnyStr(),
name="2",
path=("__pregel_push", 0),
error=None,
interrupts=(),
state=None,
result=["2|3"],
),
PregelTask(
id=AnyStr(),
name="flaky",
path=("__pregel_push", 1),
error=None,
interrupts=(
Interrupt(
value="Bahh", resumable=False, ns=None, when="during"
),
),
state=None,
result=["flaky|4"],
),
PregelTask(
id=AnyStr(),
name="3",
path=("__pregel_pull", "3"),
error=None,
interrupts=(),
state=None,
result=["3"],
),
),
),
StateSnapshot(
values=["0", "1"],
next=("2", "2", "3.1"),
config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"1": ["1"]},
"thread_id": "1",
"step": 1,
"parents": {},
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
},
tasks=(
PregelTask(
id=AnyStr(),
name="2",
path=("__pregel_push", 0),
error=None,
interrupts=(),
state=None,
result=["2|Command(goto=Send(node='2', arg=3))"],
),
PregelTask(
id=AnyStr(),
name="2",
path=("__pregel_push", 1),
error=None,
interrupts=(),
state=None,
result=["2|Command(goto=Send(node='flaky', arg=4))"],
),
PregelTask(
id=AnyStr(),
name="3.1",
@@ -2812,7 +2910,7 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
),
StateSnapshot(
values=["0"],
next=("1", "2", "2", "2", "flaky"),
next=("1",),
config={
"configurable": {
"thread_id": "1",
@@ -2845,50 +2943,6 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None:
state=None,
result=["1"],
),
PregelTask(
id=AnyStr(),
name="2",
path=("__pregel_push", ("__pregel_pull", "1"), 2),
error=None,
interrupts=(),
state=None,
result=["2|Command(goto=Send(node='2', arg=3))"],
),
PregelTask(
id=AnyStr(),
name="2",
path=("__pregel_push", ("__pregel_pull", "1"), 3),
error=None,
interrupts=(),
state=None,
result=["2|Command(goto=Send(node='flaky', arg=4))"],
),
PregelTask(
id=AnyStr(),
name="2",
path=(
"__pregel_push",
("__pregel_push", ("__pregel_pull", "1"), 2),
2,
),
error=None,
interrupts=(),
state=None,
result=["2|3"],
),
PregelTask(
id=AnyStr(),
name="flaky",
path=(
"__pregel_push",
("__pregel_push", ("__pregel_pull", "1"), 3),
2,
),
error=None,
interrupts=(Interrupt(value="Bahh", when="during"),),
state=None,
result=["flaky|4"],
),
),
),
StateSnapshot(
@@ -3046,9 +3100,6 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
}
assert foo_called == 0
if not FF_SEND_V2:
return
# get state should show the pending task
state = await graph.aget_state(thread1)
assert state == StateSnapshot(
@@ -3077,9 +3128,24 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
}
},
metadata={
"step": 0,
"step": 1,
"source": "loop",
"writes": None,
"writes": {
"agent": {
"messages": AIMessage(
content="",
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
)
}
},
"parents": {},
"thread_id": "2",
},
@@ -3096,34 +3162,10 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
}
),
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
"messages": AIMessage(
content="",
additional_kwargs={},
response_metadata={},
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
)
},
),
PregelTask(
id=AnyStr(),
name="foo",
path=("__pregel_push", ("__pregel_pull", "agent"), 2),
path=("__pregel_push", 0),
error=None,
interrupts=(),
state=None,
@@ -3157,7 +3199,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
}
},
metadata={
"step": 1,
"step": 2,
"source": "update",
"writes": {
"agent": {
@@ -3244,9 +3286,24 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
}
},
metadata={
"step": 0,
"step": 1,
"source": "loop",
"writes": None,
"writes": {
"agent": {
"messages": AIMessage(
content="",
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
)
}
},
"parents": {},
"thread_id": "3",
},
@@ -3263,32 +3320,10 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
}
),
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
"messages": AIMessage(
"",
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
)
},
),
PregelTask(
id=AnyStr(),
name="foo",
path=("__pregel_push", ("__pregel_pull", "agent"), 2),
path=("__pregel_push", 0),
error=None,
interrupts=(),
state=None,
@@ -3343,7 +3378,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
}
},
metadata={
"step": 1,
"step": 2,
"source": "update",
"writes": {
"agent": {
@@ -3379,7 +3414,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None:
PregelTask(
id=AnyStr(),
name="foo",
path=("__pregel_push", (), 0),
path=("__pregel_push", 0),
error=None,
interrupts=(),
state=None,
@@ -3530,9 +3565,6 @@ async def test_send_react_interrupt_control(
}
assert foo_called == 0
if not FF_SEND_V2:
return
# get state should show the pending task
state = await graph.aget_state(thread1)
assert state == StateSnapshot(
@@ -3561,9 +3593,24 @@ async def test_send_react_interrupt_control(
}
},
metadata={
"step": 0,
"step": 1,
"source": "loop",
"writes": None,
"writes": {
"agent": {
"messages": AIMessage(
content="",
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
)
}
},
"parents": {},
"thread_id": "2",
},
@@ -3580,34 +3627,10 @@ async def test_send_react_interrupt_control(
}
),
tasks=(
PregelTask(
id=AnyStr(),
name="agent",
path=("__pregel_pull", "agent"),
error=None,
interrupts=(),
state=None,
result={
"messages": AIMessage(
content="",
additional_kwargs={},
response_metadata={},
id="ai1",
tool_calls=[
{
"name": "foo",
"args": {"hi": [1, 2, 3]},
"id": "",
"type": "tool_call",
}
],
)
},
),
PregelTask(
id=AnyStr(),
name="foo",
path=("__pregel_push", ("__pregel_pull", "agent"), 2),
path=("__pregel_push", 0),
error=None,
interrupts=(),
state=None,
@@ -3641,7 +3664,7 @@ async def test_send_react_interrupt_control(
}
},
metadata={
"step": 1,
"step": 2,
"source": "update",
"writes": {
"agent": {
@@ -5732,9 +5755,10 @@ async def test_store_injected_async(checkpointer_name: str, store_name: str) ->
builder.add_node(f"node_{i}", Node(i))
builder.add_edge("__start__", f"node_{i}")
async with awith_checkpointer(checkpointer_name) as checkpointer, awith_store(
store_name
) as the_store:
async with (
awith_checkpointer(checkpointer_name) as checkpointer,
awith_store(store_name) as the_store,
):
graph = builder.compile(store=the_store, checkpointer=checkpointer)
# Test batch operations with multiple threads
+2 -4
View File
@@ -9,7 +9,7 @@ import pytest
from aiokafka import AIOKafkaProducer
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.constants import FF_SEND_V2, START
from langgraph.constants import START
from langgraph.errors import NodeInterrupt
from langgraph.graph.state import CompiledStateGraph, StateGraph
from langgraph.scheduler.kafka import serde
@@ -76,10 +76,8 @@ def mk_push_graph(
return builder.compile(checkpointer=checkpointer)
@pytest.mark.skip("TODO: re-enable in next PR")
async def test_push_graph(topics: Topics, acheckpointer: BaseCheckpointSaver) -> None:
if not FF_SEND_V2:
pytest.skip("Test requires FF_SEND_V2")
input = ["0"]
config = {"configurable": {"thread_id": "1"}}
graph = mk_push_graph(acheckpointer)
+2 -4
View File
@@ -8,7 +8,7 @@ from typing import (
import pytest
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.constants import FF_SEND_V2, START
from langgraph.constants import START
from langgraph.errors import NodeInterrupt
from langgraph.graph.state import CompiledStateGraph, StateGraph
from langgraph.scheduler.kafka import serde
@@ -76,10 +76,8 @@ def mk_push_graph(
return builder.compile(checkpointer=checkpointer)
@pytest.mark.skip("TODO: re-enable in next PR")
def test_push_graph(topics: Topics, acheckpointer: BaseCheckpointSaver) -> None:
if not FF_SEND_V2:
pytest.skip("Test requires FF_SEND_V2")
input = ["0"]
config = {"configurable": {"thread_id": "1"}}
graph = mk_push_graph(acheckpointer)