mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-28 18:59:42 +02:00
Remove send v2
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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": {
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user