mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-09 19:27:54 +02:00
Compare commits
@@ -248,6 +248,7 @@ class PregelLoop:
|
|||||||
self.retry_policy = retry_policy
|
self.retry_policy = retry_policy
|
||||||
self.cache_policy = cache_policy
|
self.cache_policy = cache_policy
|
||||||
self.durability = durability
|
self.durability = durability
|
||||||
|
self.skipped_task_ids: set[str] = set()
|
||||||
if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]:
|
if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]:
|
||||||
self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM])
|
self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM])
|
||||||
scratchpad: PregelScratchpad | None = config[CONF].get(CONFIG_KEY_SCRATCHPAD)
|
scratchpad: PregelScratchpad | None = config[CONF].get(CONFIG_KEY_SCRATCHPAD)
|
||||||
@@ -315,15 +316,35 @@ class PregelLoop:
|
|||||||
]
|
]
|
||||||
writes_to_save: WritesT = [
|
writes_to_save: WritesT = [
|
||||||
w[1:] for w in self.checkpoint_pending_writes if w[0] == task_id
|
w[1:] for w in self.checkpoint_pending_writes if w[0] == task_id
|
||||||
] + list(writes)
|
] + [(c, v) for c, v in writes if c != RESUME]
|
||||||
|
self.checkpoint_pending_writes.extend((task_id, c, v) for c, v in writes)
|
||||||
else:
|
else:
|
||||||
# remove existing writes for this task
|
# build map of existing interrupts for this task: interrupt id -> list of interrupts
|
||||||
|
existing_interrupts_by_id: dict[str, list[Any]] = {
|
||||||
|
v[0].id: v
|
||||||
|
for tid, ch, v in self.checkpoint_pending_writes
|
||||||
|
if tid == task_id and ch == INTERRUPT
|
||||||
|
}
|
||||||
|
writes_to_save = []
|
||||||
|
for ch, v in writes:
|
||||||
|
if ch == INTERRUPT:
|
||||||
|
# we merge new interrupt writes with existing interrupts writes if they
|
||||||
|
# occurred within the same task (which means they have the same interrupt id)
|
||||||
|
new_interrupts = v if isinstance(v, list) else list(v)
|
||||||
|
if new_interrupts and (
|
||||||
|
existing := existing_interrupts_by_id.get(new_interrupts[0].id)
|
||||||
|
):
|
||||||
|
v = existing + new_interrupts
|
||||||
|
writes_to_save.append((ch, v))
|
||||||
|
else:
|
||||||
|
# we add non-interrupt writes as-is
|
||||||
|
writes_to_save.append((ch, v))
|
||||||
|
|
||||||
|
# replace all writes for this task_id with the merged writes
|
||||||
self.checkpoint_pending_writes = [
|
self.checkpoint_pending_writes = [
|
||||||
w for w in self.checkpoint_pending_writes if w[0] != task_id
|
w for w in self.checkpoint_pending_writes if w[0] != task_id
|
||||||
]
|
] + [(task_id, c, v) for c, v in writes_to_save]
|
||||||
writes_to_save = writes
|
|
||||||
# save writes
|
|
||||||
self.checkpoint_pending_writes.extend((task_id, c, v) for c, v in writes)
|
|
||||||
if self.durability != "exit" and self.checkpointer_put_writes is not None:
|
if self.durability != "exit" and self.checkpointer_put_writes is not None:
|
||||||
config = patch_configurable(
|
config = patch_configurable(
|
||||||
self.checkpoint_config,
|
self.checkpoint_config,
|
||||||
@@ -471,6 +492,22 @@ class PregelLoop:
|
|||||||
cache_policy=self.cache_policy,
|
cache_policy=self.cache_policy,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
resume_map = self.config.get(CONF, {}).get(CONFIG_KEY_RESUME_MAP, {})
|
||||||
|
if resume_map or self.input is None:
|
||||||
|
# do not re-execute tasks that have unresumable interrupts
|
||||||
|
# i.e. when the graph is invoked with None, or the interrupt id is not in the resume map
|
||||||
|
skipped_interrupt_ids = self._pending_interrupts() - set(resume_map)
|
||||||
|
self.skipped_task_ids = {
|
||||||
|
task_id
|
||||||
|
for task_id, channel, value in self.checkpoint_pending_writes
|
||||||
|
if channel == INTERRUPT
|
||||||
|
# interrupts within a task are uncovered sequentially as resumes are provided,
|
||||||
|
# so we only need to check the last interrupt id
|
||||||
|
and value[-1].id in skipped_interrupt_ids
|
||||||
|
}
|
||||||
|
else:
|
||||||
|
self.skipped_task_ids = set()
|
||||||
|
|
||||||
# produce debug output
|
# produce debug output
|
||||||
if self._checkpointer_put_after_previous is not None:
|
if self._checkpointer_put_after_previous is not None:
|
||||||
self._emit(
|
self._emit(
|
||||||
@@ -516,9 +553,45 @@ class PregelLoop:
|
|||||||
if task.writes:
|
if task.writes:
|
||||||
self.output_writes(task.id, task.writes, cached=True)
|
self.output_writes(task.id, task.writes, cached=True)
|
||||||
|
|
||||||
|
if self.skipped_task_ids:
|
||||||
|
# remove tasks with writes that have been matched with previous pending writes
|
||||||
|
self.skipped_task_ids = {
|
||||||
|
task_id
|
||||||
|
for task_id in self.skipped_task_ids
|
||||||
|
if not self.tasks[task_id].writes
|
||||||
|
}
|
||||||
|
# output interrupt writes for blocked tasks so they are still visible in the stream
|
||||||
|
for task_id, channel, value in self.checkpoint_pending_writes:
|
||||||
|
if task_id in self.skipped_task_ids and channel == INTERRUPT:
|
||||||
|
# find resume count for this task
|
||||||
|
resumes = next(
|
||||||
|
(
|
||||||
|
v
|
||||||
|
for tid, ch, v in self.checkpoint_pending_writes
|
||||||
|
if tid == task_id and ch == RESUME
|
||||||
|
),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
resume_count = len(resumes) if resumes is not None else 0
|
||||||
|
# only output unresumed interrupts
|
||||||
|
if resume_count < len(value):
|
||||||
|
self.output_writes(task_id, [(INTERRUPT, value[resume_count:])])
|
||||||
|
|
||||||
return True
|
return True
|
||||||
|
|
||||||
def after_tick(self) -> None:
|
def after_tick(self) -> None:
|
||||||
|
if self.skipped_task_ids:
|
||||||
|
# raise early GraphInterrupt for skipped tasks.
|
||||||
|
# since we know len(resumes) < len(interrupts) for these tasks, we
|
||||||
|
# can prevent unnecessary node re-execution by raising early
|
||||||
|
interrupts = []
|
||||||
|
for task_id, channel, value in self.checkpoint_pending_writes:
|
||||||
|
if channel == INTERRUPT and task_id in self.skipped_task_ids:
|
||||||
|
interrupts.extend(value)
|
||||||
|
if interrupts:
|
||||||
|
raise GraphInterrupt(interrupts)
|
||||||
|
|
||||||
|
self.skipped_task_ids.clear()
|
||||||
# finish superstep
|
# finish superstep
|
||||||
writes = [w for t in self.tasks.values() for w in t.writes]
|
writes = [w for t in self.tasks.values() for w in t.writes]
|
||||||
# all tasks have finished
|
# all tasks have finished
|
||||||
@@ -570,30 +643,26 @@ class PregelLoop:
|
|||||||
|
|
||||||
def _pending_interrupts(self) -> set[str]:
|
def _pending_interrupts(self) -> set[str]:
|
||||||
"""Return the set of interrupt ids that are pending without corresponding resume values."""
|
"""Return the set of interrupt ids that are pending without corresponding resume values."""
|
||||||
# mapping of task ids to interrupt ids
|
# mapping of task ids to (interrupt_id, interrupt_count)
|
||||||
pending_interrupts: dict[str, str] = {}
|
pending_interrupts: dict[str, tuple[str, int]] = {}
|
||||||
|
# mapping of task ids to resume count
|
||||||
|
pending_resumes: dict[str, int] = {}
|
||||||
|
|
||||||
# set of resume task ids
|
for task_id, channel, value in self.checkpoint_pending_writes:
|
||||||
pending_resumes: set[str] = set()
|
if channel == INTERRUPT:
|
||||||
|
pending_interrupts[task_id] = (
|
||||||
|
value[0].id,
|
||||||
|
len(value),
|
||||||
|
)
|
||||||
|
elif channel == RESUME:
|
||||||
|
resume_list = value if isinstance(value, list) else [value]
|
||||||
|
pending_resumes[task_id] = len(resume_list)
|
||||||
|
|
||||||
for task_id, write_type, value in self.checkpoint_pending_writes:
|
# keep only interrupt ids where resume_count < interrupt_count
|
||||||
if write_type == INTERRUPT:
|
|
||||||
# interrupts is always a list, but there should only be one element
|
|
||||||
pending_interrupts[task_id] = value[0].id
|
|
||||||
elif write_type == RESUME:
|
|
||||||
pending_resumes.add(task_id)
|
|
||||||
|
|
||||||
resumed_interrupt_ids = {
|
|
||||||
pending_interrupts[task_id]
|
|
||||||
for task_id in pending_resumes
|
|
||||||
if task_id in pending_interrupts
|
|
||||||
}
|
|
||||||
|
|
||||||
# Keep only interrupts whose interrupt_id is not resumed
|
|
||||||
hanging_interrupts: set[str] = {
|
hanging_interrupts: set[str] = {
|
||||||
interrupt_id
|
interrupt_id
|
||||||
for interrupt_id in pending_interrupts.values()
|
for task_id, (interrupt_id, interrupt_count) in pending_interrupts.items()
|
||||||
if interrupt_id not in resumed_interrupt_ids
|
if pending_resumes.get(task_id, 0) < interrupt_count
|
||||||
}
|
}
|
||||||
|
|
||||||
return hanging_interrupts
|
return hanging_interrupts
|
||||||
@@ -1026,6 +1095,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
|
|||||||
|
|
||||||
def put_writes(self, task_id: str, writes: WritesT) -> None:
|
def put_writes(self, task_id: str, writes: WritesT) -> None:
|
||||||
"""Put writes for a task, to be read by the next tick."""
|
"""Put writes for a task, to be read by the next tick."""
|
||||||
|
|
||||||
super().put_writes(task_id, writes)
|
super().put_writes(task_id, writes)
|
||||||
if not writes or self.cache is None or not hasattr(self, "tasks"):
|
if not writes or self.cache is None or not hasattr(self, "tasks"):
|
||||||
return
|
return
|
||||||
|
|||||||
@@ -2672,7 +2672,11 @@ class Pregel(
|
|||||||
for task in loop.match_cached_writes():
|
for task in loop.match_cached_writes():
|
||||||
loop.output_writes(task.id, task.writes, cached=True)
|
loop.output_writes(task.id, task.writes, cached=True)
|
||||||
for _ in runner.tick(
|
for _ in runner.tick(
|
||||||
[t for t in loop.tasks.values() if not t.writes],
|
[
|
||||||
|
t
|
||||||
|
for t in loop.tasks.values()
|
||||||
|
if not t.writes and t.id not in loop.skipped_task_ids
|
||||||
|
],
|
||||||
timeout=self.step_timeout,
|
timeout=self.step_timeout,
|
||||||
get_waiter=get_waiter,
|
get_waiter=get_waiter,
|
||||||
schedule_task=loop.accept_push,
|
schedule_task=loop.accept_push,
|
||||||
@@ -2991,7 +2995,11 @@ class Pregel(
|
|||||||
for task in await loop.amatch_cached_writes():
|
for task in await loop.amatch_cached_writes():
|
||||||
loop.output_writes(task.id, task.writes, cached=True)
|
loop.output_writes(task.id, task.writes, cached=True)
|
||||||
async for _ in runner.atick(
|
async for _ in runner.atick(
|
||||||
[t for t in loop.tasks.values() if not t.writes],
|
[
|
||||||
|
t
|
||||||
|
for t in loop.tasks.values()
|
||||||
|
if not t.writes and t.id not in loop.skipped_task_ids
|
||||||
|
],
|
||||||
timeout=self.step_timeout,
|
timeout=self.step_timeout,
|
||||||
get_waiter=get_waiter,
|
get_waiter=get_waiter,
|
||||||
schedule_task=loop.aaccept_push,
|
schedule_task=loop.aaccept_push,
|
||||||
|
|||||||
@@ -1,12 +1,21 @@
|
|||||||
|
import operator
|
||||||
|
import sys
|
||||||
|
from typing import Annotated
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
from langgraph.checkpoint.base import BaseCheckpointSaver
|
from langgraph.checkpoint.base import BaseCheckpointSaver
|
||||||
from typing_extensions import TypedDict
|
from typing_extensions import TypedDict
|
||||||
|
|
||||||
from langgraph.graph import END, START, StateGraph
|
from langgraph.graph import END, START, StateGraph
|
||||||
from langgraph.types import Durability
|
from langgraph.types import Command, Durability, Send, interrupt
|
||||||
|
|
||||||
pytestmark = pytest.mark.anyio
|
pytestmark = pytest.mark.anyio
|
||||||
|
|
||||||
|
NEEDS_CONTEXTVARS = pytest.mark.skipif(
|
||||||
|
sys.version_info < (3, 11),
|
||||||
|
reason="Python 3.11+ is required for async contextvars support",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def test_interruption_without_state_updates(
|
def test_interruption_without_state_updates(
|
||||||
sync_checkpointer: BaseCheckpointSaver, durability: Durability
|
sync_checkpointer: BaseCheckpointSaver, durability: Durability
|
||||||
@@ -90,3 +99,567 @@ async def test_interruption_without_state_updates_async(
|
|||||||
assert (await graph.aget_state(thread)).next == ()
|
assert (await graph.aget_state(thread)).next == ()
|
||||||
n_checkpoints = len([c async for c in graph.aget_state_history(thread)])
|
n_checkpoints = len([c async for c in graph.aget_state_history(thread)])
|
||||||
assert n_checkpoints == (5 if durability != "exit" else 3)
|
assert n_checkpoints == (5 if durability != "exit" else 3)
|
||||||
|
|
||||||
|
|
||||||
|
def test_interrupt_with_send_payloads(sync_checkpointer: BaseCheckpointSaver) -> None:
|
||||||
|
"""Test interruption in map node with Send payloads and human-in-the-loop resume."""
|
||||||
|
|
||||||
|
# Global counter to track node executions
|
||||||
|
node_counter = {"entry": 0, "map_node": 0}
|
||||||
|
|
||||||
|
class State(TypedDict):
|
||||||
|
items: list[str]
|
||||||
|
processed: Annotated[list[str], operator.add]
|
||||||
|
|
||||||
|
def entry_node(state: State):
|
||||||
|
node_counter["entry"] += 1
|
||||||
|
return {} # No state updates in entry node
|
||||||
|
|
||||||
|
def send_to_map(state: State):
|
||||||
|
return [Send("map_node", {"item": item}) for item in state["items"]]
|
||||||
|
|
||||||
|
def map_node(state: State):
|
||||||
|
node_counter["map_node"] += 1
|
||||||
|
if "dangerous" in state["item"]:
|
||||||
|
value = interrupt({"processing": state["item"]})
|
||||||
|
return {"processed": [f"processed_{value}"]}
|
||||||
|
else:
|
||||||
|
return {"processed": [f"processed_{state['item']}_auto"]}
|
||||||
|
|
||||||
|
builder = StateGraph(State)
|
||||||
|
builder.add_node("entry", entry_node)
|
||||||
|
builder.add_node("map_node", map_node)
|
||||||
|
builder.add_edge(START, "entry")
|
||||||
|
builder.add_conditional_edges("entry", send_to_map, ["map_node"])
|
||||||
|
builder.add_edge("map_node", END)
|
||||||
|
|
||||||
|
graph = builder.compile(checkpointer=sync_checkpointer)
|
||||||
|
|
||||||
|
config = {"configurable": {"thread_id": "test_interrupt_send"}}
|
||||||
|
|
||||||
|
# Run until interrupts
|
||||||
|
result = graph.invoke(
|
||||||
|
{"items": ["item1", "dangerous_item1", "dangerous_item2"]}, config=config
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify we have interrupts (only one for dangerous_item)
|
||||||
|
interrupts = result.get("__interrupt__", [])
|
||||||
|
assert len(interrupts) == 2
|
||||||
|
assert "dangerous_item" in interrupts[0].value["processing"]
|
||||||
|
|
||||||
|
# Resume with mapping of interrupt IDs to values
|
||||||
|
resume_map = {i.id: f"human_input_{i.value['processing']}" for i in interrupts}
|
||||||
|
|
||||||
|
final_result = graph.invoke(Command(resume=resume_map), config=config)
|
||||||
|
|
||||||
|
# Verify final result contains processed items
|
||||||
|
assert "processed" in final_result
|
||||||
|
processed_items = final_result["processed"]
|
||||||
|
assert len(processed_items) == 3
|
||||||
|
assert "processed_item1_auto" in processed_items # item1 processed automatically
|
||||||
|
assert any(
|
||||||
|
"processed_human_input_dangerous_item1" in item for item in processed_items
|
||||||
|
) # dangerous_item1 processed after interrupt
|
||||||
|
assert any(
|
||||||
|
"processed_human_input_dangerous_item2" in item for item in processed_items
|
||||||
|
) # dangerous_item2 processed after interrupt
|
||||||
|
|
||||||
|
# Verify node execution counts
|
||||||
|
assert node_counter["entry"] == 1 # Entry node runs once
|
||||||
|
# Map node runs 3 times initially (item1 completes, 2 dangerous_items interrupt),
|
||||||
|
# then 2 times on resume
|
||||||
|
assert node_counter["map_node"] == 5
|
||||||
|
|
||||||
|
|
||||||
|
@NEEDS_CONTEXTVARS
|
||||||
|
async def test_interrupt_with_send_payloads_async(
|
||||||
|
async_checkpointer: BaseCheckpointSaver, durability: Durability
|
||||||
|
) -> None:
|
||||||
|
"""Test interruption in map node with Send payloads and human-in-the-loop resume."""
|
||||||
|
|
||||||
|
# Global counter to track node executions
|
||||||
|
node_counter = {"entry": 0, "map_node": 0}
|
||||||
|
|
||||||
|
class State(TypedDict):
|
||||||
|
items: list[str]
|
||||||
|
processed: Annotated[list[str], operator.add]
|
||||||
|
|
||||||
|
def entry_node(state: State):
|
||||||
|
node_counter["entry"] += 1
|
||||||
|
return {} # No state updates in entry node
|
||||||
|
|
||||||
|
def send_to_map(state: State):
|
||||||
|
return [Send("map_node", {"item": item}) for item in state["items"]]
|
||||||
|
|
||||||
|
def map_node(state: State):
|
||||||
|
node_counter["map_node"] += 1
|
||||||
|
if "dangerous" in state["item"]:
|
||||||
|
value = interrupt({"processing": state["item"]})
|
||||||
|
return {"processed": [f"processed_{value}"]}
|
||||||
|
else:
|
||||||
|
return {"processed": [f"processed_{state['item']}_auto"]}
|
||||||
|
|
||||||
|
builder = StateGraph(State)
|
||||||
|
builder.add_node("entry", entry_node)
|
||||||
|
builder.add_node("map_node", map_node)
|
||||||
|
builder.add_edge(START, "entry")
|
||||||
|
builder.add_conditional_edges("entry", send_to_map, ["map_node"])
|
||||||
|
builder.add_edge("map_node", END)
|
||||||
|
|
||||||
|
graph = builder.compile(checkpointer=async_checkpointer)
|
||||||
|
|
||||||
|
config = {"configurable": {"thread_id": "test_interrupt_send"}}
|
||||||
|
|
||||||
|
# Run until interrupts
|
||||||
|
result = await graph.ainvoke(
|
||||||
|
{"items": ["item1", "dangerous_item1", "dangerous_item2"]}, config=config
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify we have interrupts (only one for dangerous_item)
|
||||||
|
interrupts = result.get("__interrupt__", [])
|
||||||
|
assert len(interrupts) == 2
|
||||||
|
assert "dangerous_item" in interrupts[0].value["processing"]
|
||||||
|
|
||||||
|
# Resume with mapping of interrupt IDs to values
|
||||||
|
resume_map = {i.id: f"human_input_{i.value['processing']}" for i in interrupts}
|
||||||
|
|
||||||
|
final_result = await graph.ainvoke(Command(resume=resume_map), config=config)
|
||||||
|
|
||||||
|
# Verify final result contains processed items
|
||||||
|
assert "processed" in final_result
|
||||||
|
processed_items = final_result["processed"]
|
||||||
|
assert len(processed_items) == 3
|
||||||
|
assert "processed_item1_auto" in processed_items # item1 processed automatically
|
||||||
|
assert any(
|
||||||
|
"processed_human_input_dangerous_item1" in item for item in processed_items
|
||||||
|
) # dangerous_item1 processed after interrupt
|
||||||
|
assert any(
|
||||||
|
"processed_human_input_dangerous_item2" in item for item in processed_items
|
||||||
|
) # dangerous_item2 processed after interrupt
|
||||||
|
|
||||||
|
# Verify node execution counts
|
||||||
|
assert node_counter["entry"] == 1 # Entry node runs once
|
||||||
|
# Map node runs 3 times initially (item1 completes, 2 dangerous_items interrupt),
|
||||||
|
# then 2 times on resume
|
||||||
|
assert node_counter["map_node"] == 5
|
||||||
|
|
||||||
|
|
||||||
|
def test_interrupt_with_send_payloads_sequential_resume(
|
||||||
|
sync_checkpointer: BaseCheckpointSaver,
|
||||||
|
) -> None:
|
||||||
|
"""Test interruption in map node with Send payloads and sequential resume."""
|
||||||
|
|
||||||
|
# Global counter to track node executions
|
||||||
|
node_counter = {"entry": 0, "map_node": 0}
|
||||||
|
|
||||||
|
class State(TypedDict):
|
||||||
|
items: list[str]
|
||||||
|
processed: Annotated[list[str], operator.add]
|
||||||
|
|
||||||
|
def entry_node(state: State):
|
||||||
|
node_counter["entry"] += 1
|
||||||
|
return {} # No state updates in entry node
|
||||||
|
|
||||||
|
def send_to_map(state: State):
|
||||||
|
return [Send("map_node", {"item": item}) for item in state["items"]]
|
||||||
|
|
||||||
|
def map_node(state: State):
|
||||||
|
node_counter["map_node"] += 1
|
||||||
|
if "dangerous" in state["item"]:
|
||||||
|
value = interrupt({"processing": state["item"]})
|
||||||
|
return {"processed": [f"processed_{value}"]}
|
||||||
|
else:
|
||||||
|
return {"processed": [f"processed_{state['item']}_auto"]}
|
||||||
|
|
||||||
|
builder = StateGraph(State)
|
||||||
|
builder.add_node("entry", entry_node)
|
||||||
|
builder.add_node("map_node", map_node)
|
||||||
|
builder.add_edge(START, "entry")
|
||||||
|
builder.add_conditional_edges("entry", send_to_map, ["map_node"])
|
||||||
|
builder.add_edge("map_node", END)
|
||||||
|
|
||||||
|
graph = builder.compile(checkpointer=sync_checkpointer)
|
||||||
|
|
||||||
|
config = {"configurable": {"thread_id": "test_interrupt_send_sequential"}}
|
||||||
|
|
||||||
|
# Run until interrupts
|
||||||
|
result = graph.invoke(
|
||||||
|
{"items": ["item1", "dangerous_item1", "dangerous_item2"]}, config=config
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify we have interrupts
|
||||||
|
interrupts = result.get("__interrupt__", [])
|
||||||
|
assert len(interrupts) == 2
|
||||||
|
assert "dangerous_item" in interrupts[0].value["processing"]
|
||||||
|
|
||||||
|
# Resume first interrupt only
|
||||||
|
first_interrupt = interrupts[0]
|
||||||
|
first_resume_map = {
|
||||||
|
first_interrupt.id: f"human_input_{first_interrupt.value['processing']}"
|
||||||
|
}
|
||||||
|
|
||||||
|
partial_result = graph.invoke(Command(resume=first_resume_map), config=config)
|
||||||
|
|
||||||
|
# Verify we still have one pending interrupt
|
||||||
|
remaining_interrupts = partial_result.get("__interrupt__", [])
|
||||||
|
assert len(remaining_interrupts) == 1
|
||||||
|
|
||||||
|
# Resume second interrupt
|
||||||
|
second_interrupt = remaining_interrupts[0]
|
||||||
|
second_resume_map = {
|
||||||
|
second_interrupt.id: f"human_input_{second_interrupt.value['processing']}"
|
||||||
|
}
|
||||||
|
|
||||||
|
final_result = graph.invoke(Command(resume=second_resume_map), config=config)
|
||||||
|
|
||||||
|
# Verify final result contains processed items
|
||||||
|
assert "processed" in final_result
|
||||||
|
processed_items = final_result["processed"]
|
||||||
|
assert len(processed_items) == 3
|
||||||
|
assert "processed_item1_auto" in processed_items # item1 processed automatically
|
||||||
|
assert any(
|
||||||
|
"processed_human_input_dangerous_item1" in item for item in processed_items
|
||||||
|
) # dangerous_item1 processed after interrupt
|
||||||
|
assert any(
|
||||||
|
"processed_human_input_dangerous_item2" in item for item in processed_items
|
||||||
|
) # dangerous_item2 processed after interrupt
|
||||||
|
|
||||||
|
# Verify node execution counts
|
||||||
|
assert node_counter["entry"] == 1 # Entry node runs once
|
||||||
|
# Map node runs 3 times initially (item1 completes, 2 dangerous_items interrupt),
|
||||||
|
# then 1 time on first resume, then 1 time on second resume
|
||||||
|
assert node_counter["map_node"] == 5
|
||||||
|
|
||||||
|
|
||||||
|
@NEEDS_CONTEXTVARS
|
||||||
|
async def test_interrupt_with_send_payloads_sequential_resume_async(
|
||||||
|
async_checkpointer: BaseCheckpointSaver,
|
||||||
|
) -> None:
|
||||||
|
"""Test interruption in map node with Send payloads and sequential resume."""
|
||||||
|
|
||||||
|
# Global counter to track node executions
|
||||||
|
node_counter = {"entry": 0, "map_node": 0}
|
||||||
|
|
||||||
|
class State(TypedDict):
|
||||||
|
items: list[str]
|
||||||
|
processed: Annotated[list[str], operator.add]
|
||||||
|
|
||||||
|
def entry_node(state: State):
|
||||||
|
node_counter["entry"] += 1
|
||||||
|
return {} # No state updates in entry node
|
||||||
|
|
||||||
|
def send_to_map(state: State):
|
||||||
|
return [Send("map_node", {"item": item}) for item in state["items"]]
|
||||||
|
|
||||||
|
def map_node(state: State):
|
||||||
|
node_counter["map_node"] += 1
|
||||||
|
if "dangerous" in state["item"]:
|
||||||
|
value = interrupt({"processing": state["item"]})
|
||||||
|
return {"processed": [f"processed_{value}"]}
|
||||||
|
else:
|
||||||
|
return {"processed": [f"processed_{state['item']}_auto"]}
|
||||||
|
|
||||||
|
builder = StateGraph(State)
|
||||||
|
builder.add_node("entry", entry_node)
|
||||||
|
builder.add_node("map_node", map_node)
|
||||||
|
builder.add_edge(START, "entry")
|
||||||
|
builder.add_conditional_edges("entry", send_to_map, ["map_node"])
|
||||||
|
builder.add_edge("map_node", END)
|
||||||
|
|
||||||
|
graph = builder.compile(checkpointer=async_checkpointer)
|
||||||
|
|
||||||
|
config = {"configurable": {"thread_id": "test_interrupt_send_sequential"}}
|
||||||
|
|
||||||
|
# Run until interrupts
|
||||||
|
result = await graph.ainvoke(
|
||||||
|
{"items": ["item1", "dangerous_item1", "dangerous_item2"]}, config=config
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify we have interrupts
|
||||||
|
interrupts = result.get("__interrupt__", [])
|
||||||
|
assert len(interrupts) == 2
|
||||||
|
assert "dangerous_item" in interrupts[0].value["processing"]
|
||||||
|
|
||||||
|
# Resume first interrupt only
|
||||||
|
first_interrupt = interrupts[0]
|
||||||
|
first_resume_map = {
|
||||||
|
first_interrupt.id: f"human_input_{first_interrupt.value['processing']}"
|
||||||
|
}
|
||||||
|
|
||||||
|
partial_result = await graph.ainvoke(
|
||||||
|
Command(resume=first_resume_map), config=config
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify we still have one pending interrupt
|
||||||
|
remaining_interrupts = partial_result.get("__interrupt__", [])
|
||||||
|
assert len(remaining_interrupts) == 1
|
||||||
|
|
||||||
|
# Resume second interrupt
|
||||||
|
second_interrupt = remaining_interrupts[0]
|
||||||
|
second_resume_map = {
|
||||||
|
second_interrupt.id: f"human_input_{second_interrupt.value['processing']}"
|
||||||
|
}
|
||||||
|
|
||||||
|
final_result = await graph.ainvoke(Command(resume=second_resume_map), config=config)
|
||||||
|
|
||||||
|
# Verify final result contains processed items
|
||||||
|
assert "processed" in final_result
|
||||||
|
processed_items = final_result["processed"]
|
||||||
|
assert len(processed_items) == 3
|
||||||
|
assert "processed_item1_auto" in processed_items # item1 processed automatically
|
||||||
|
assert any(
|
||||||
|
"processed_human_input_dangerous_item1" in item for item in processed_items
|
||||||
|
) # dangerous_item1 processed after interrupt
|
||||||
|
assert any(
|
||||||
|
"processed_human_input_dangerous_item2" in item for item in processed_items
|
||||||
|
) # dangerous_item2 processed after interrupt
|
||||||
|
|
||||||
|
# Verify node execution counts
|
||||||
|
assert node_counter["entry"] == 1 # Entry node runs once
|
||||||
|
# Map node runs 3 times initially (item1 completes, 2 dangerous_items interrupt),
|
||||||
|
# then 1 time on first resume, then 1 time on second resume
|
||||||
|
assert node_counter["map_node"] == 5
|
||||||
|
|
||||||
|
|
||||||
|
def test_node_with_multiple_interrupts_requires_full_resume(
|
||||||
|
sync_checkpointer: BaseCheckpointSaver,
|
||||||
|
) -> None:
|
||||||
|
"""Test a number of different resume patterns for a node with multiple interrupts,
|
||||||
|
|
||||||
|
Ensures that a node is not re-executed until valid resume values have been provided to all
|
||||||
|
discovered interrupts"""
|
||||||
|
|
||||||
|
node_counter = 0
|
||||||
|
|
||||||
|
class State(TypedDict):
|
||||||
|
input: str
|
||||||
|
|
||||||
|
def double_interrupt_node(state: State):
|
||||||
|
nonlocal node_counter
|
||||||
|
node_counter += 1
|
||||||
|
first = interrupt("first")
|
||||||
|
second = interrupt("second")
|
||||||
|
third = interrupt("third")
|
||||||
|
return {"input": f"{first}-{second}-{third}"}
|
||||||
|
|
||||||
|
builder = StateGraph(State)
|
||||||
|
builder.add_node("double_interrupt", double_interrupt_node)
|
||||||
|
builder.add_edge(START, "double_interrupt")
|
||||||
|
builder.add_edge("double_interrupt", END)
|
||||||
|
|
||||||
|
graph = builder.compile(checkpointer=sync_checkpointer)
|
||||||
|
|
||||||
|
config = {"configurable": {"thread_id": "test_double_interrupt"}}
|
||||||
|
|
||||||
|
result = graph.invoke({"input": "start"}, config=config)
|
||||||
|
|
||||||
|
interrupts = result.get("__interrupt__", [])
|
||||||
|
assert len(interrupts) == 1
|
||||||
|
first_interrupt = interrupts[0]
|
||||||
|
assert node_counter == 1
|
||||||
|
|
||||||
|
# invoke with an interrupt map that matches double_interrupt_node.
|
||||||
|
# this should execute the node
|
||||||
|
partial = graph.invoke(
|
||||||
|
Command(resume={first_interrupt.id: "human_first"}), config=config
|
||||||
|
)
|
||||||
|
remaining_interrupts = partial.get("__interrupt__", [])
|
||||||
|
assert len(remaining_interrupts) == 1
|
||||||
|
assert remaining_interrupts[0].value == "second"
|
||||||
|
assert node_counter == 2
|
||||||
|
|
||||||
|
# invoke with an interrupt map that DOES NOT match double_interrupt_node.
|
||||||
|
# this should not execute the node because the optimization kicks in
|
||||||
|
partial = graph.invoke(
|
||||||
|
Command(resume={"00000000000000000000000000000000": "nothing_burger"}),
|
||||||
|
config=config,
|
||||||
|
)
|
||||||
|
remaining_interrupts = partial.get("__interrupt__", [])
|
||||||
|
assert len(remaining_interrupts) == 1
|
||||||
|
assert remaining_interrupts[0].value == "second"
|
||||||
|
assert node_counter == 2
|
||||||
|
|
||||||
|
# invoke with None resume. this should NOT execute the node
|
||||||
|
partial = graph.invoke(None, config=config)
|
||||||
|
remaining_interrupts = partial.get("__interrupt__", [])
|
||||||
|
assert len(remaining_interrupts) == 1
|
||||||
|
assert remaining_interrupts[0].value == "second"
|
||||||
|
assert node_counter == 2
|
||||||
|
|
||||||
|
# invoke with nonspecific resume. this should execute the node
|
||||||
|
partial = graph.invoke(Command(resume="human_second"), config=config)
|
||||||
|
remaining_interrupts = partial.get("__interrupt__", [])
|
||||||
|
assert len(remaining_interrupts) == 1
|
||||||
|
print("REMAINING INTERRUPTS: ", remaining_interrupts)
|
||||||
|
assert remaining_interrupts[0].value == "third"
|
||||||
|
assert node_counter == 3
|
||||||
|
|
||||||
|
# finally, invoke with an interrupt map that matches double_interrupt_node.
|
||||||
|
# this should execute the node and all interrupts should be resolved
|
||||||
|
final_result = graph.invoke(Command(resume="human_third"), config=config)
|
||||||
|
assert "input" in final_result
|
||||||
|
assert final_result["input"] == "human_first-human_second-human_third"
|
||||||
|
assert node_counter == 4
|
||||||
|
|
||||||
|
|
||||||
|
@NEEDS_CONTEXTVARS
|
||||||
|
async def test_node_with_multiple_interrupts_requires_full_resume_async(
|
||||||
|
async_checkpointer: BaseCheckpointSaver,
|
||||||
|
) -> None:
|
||||||
|
"""Test a number of different resume patterns for a node with multiple interrupts,
|
||||||
|
|
||||||
|
Ensures that a node is not re-executed until valid resume values have been provided to all
|
||||||
|
discovered interrupts"""
|
||||||
|
|
||||||
|
node_counter = 0
|
||||||
|
|
||||||
|
class State(TypedDict):
|
||||||
|
input: str
|
||||||
|
|
||||||
|
def double_interrupt_node(state: State):
|
||||||
|
nonlocal node_counter
|
||||||
|
node_counter += 1
|
||||||
|
first = interrupt("first")
|
||||||
|
second = interrupt("second")
|
||||||
|
third = interrupt("third")
|
||||||
|
return {"input": f"{first}-{second}-{third}"}
|
||||||
|
|
||||||
|
builder = StateGraph(State)
|
||||||
|
builder.add_node("double_interrupt", double_interrupt_node)
|
||||||
|
builder.add_edge(START, "double_interrupt")
|
||||||
|
builder.add_edge("double_interrupt", END)
|
||||||
|
|
||||||
|
graph = builder.compile(checkpointer=async_checkpointer)
|
||||||
|
|
||||||
|
config = {"configurable": {"thread_id": "test_double_interrupt"}}
|
||||||
|
|
||||||
|
result = await graph.ainvoke({"input": "start"}, config=config)
|
||||||
|
|
||||||
|
interrupts = result.get("__interrupt__", [])
|
||||||
|
assert len(interrupts) == 1
|
||||||
|
first_interrupt = interrupts[0]
|
||||||
|
assert node_counter == 1
|
||||||
|
|
||||||
|
# invoke with an interrupt map that matches double_interrupt_node.
|
||||||
|
# this should execute the node
|
||||||
|
partial = await graph.ainvoke(
|
||||||
|
Command(resume={first_interrupt.id: "human_first"}), config=config
|
||||||
|
)
|
||||||
|
remaining_interrupts = partial.get("__interrupt__", [])
|
||||||
|
assert len(remaining_interrupts) == 1
|
||||||
|
assert remaining_interrupts[0].value == "second"
|
||||||
|
assert node_counter == 2
|
||||||
|
|
||||||
|
# invoke with an interrupt map that DOES NOT match double_interrupt_node.
|
||||||
|
# this should not execute the node because the optimization kicks in
|
||||||
|
partial = await graph.ainvoke(
|
||||||
|
Command(resume={"00000000000000000000000000000000": "nothing_burger"}),
|
||||||
|
config=config,
|
||||||
|
)
|
||||||
|
remaining_interrupts = partial.get("__interrupt__", [])
|
||||||
|
assert len(remaining_interrupts) == 1
|
||||||
|
assert remaining_interrupts[0].value == "second"
|
||||||
|
assert node_counter == 2
|
||||||
|
|
||||||
|
# invoke with None resume. this should NOT execute the node
|
||||||
|
partial = await graph.ainvoke(None, config=config)
|
||||||
|
remaining_interrupts = partial.get("__interrupt__", [])
|
||||||
|
assert len(remaining_interrupts) == 1
|
||||||
|
assert remaining_interrupts[0].value == "second"
|
||||||
|
assert node_counter == 2
|
||||||
|
|
||||||
|
# invoke with nonspecific resume. this should execute the node
|
||||||
|
partial = await graph.ainvoke(Command(resume="human_second"), config=config)
|
||||||
|
remaining_interrupts = partial.get("__interrupt__", [])
|
||||||
|
assert len(remaining_interrupts) == 1
|
||||||
|
print("REMAINING INTERRUPTS: ", remaining_interrupts)
|
||||||
|
assert remaining_interrupts[0].value == "third"
|
||||||
|
assert node_counter == 3
|
||||||
|
|
||||||
|
# finally, invoke with an interrupt map that matches double_interrupt_node.
|
||||||
|
# this should execute the node and all interrupts should be resolved
|
||||||
|
final_result = await graph.ainvoke(Command(resume="human_third"), config=config)
|
||||||
|
assert "input" in final_result
|
||||||
|
assert final_result["input"] == "human_first-human_second-human_third"
|
||||||
|
assert node_counter == 4
|
||||||
|
|
||||||
|
|
||||||
|
def test_invoke_interrupted_graph_with_none(
|
||||||
|
sync_checkpointer: BaseCheckpointSaver,
|
||||||
|
) -> None:
|
||||||
|
"""Test that invoking an interrupted graph with None does not duplicate interrupt writes"""
|
||||||
|
|
||||||
|
node_counter = 0
|
||||||
|
|
||||||
|
class State(TypedDict):
|
||||||
|
input: str
|
||||||
|
|
||||||
|
def double_interrupt_node(state: State):
|
||||||
|
nonlocal node_counter
|
||||||
|
node_counter += 1
|
||||||
|
first = interrupt("first")
|
||||||
|
second = interrupt("second")
|
||||||
|
return {"input": f"{first}-{second}"}
|
||||||
|
|
||||||
|
builder = StateGraph(State)
|
||||||
|
builder.add_node("double_interrupt", double_interrupt_node)
|
||||||
|
builder.add_edge(START, "double_interrupt")
|
||||||
|
builder.add_edge("double_interrupt", END)
|
||||||
|
|
||||||
|
graph = builder.compile(checkpointer=sync_checkpointer)
|
||||||
|
|
||||||
|
config = {"configurable": {"thread_id": "test_none_resume"}}
|
||||||
|
|
||||||
|
result = graph.invoke({"input": "start"}, config=config)
|
||||||
|
first_history = list(graph.get_state_history(config))
|
||||||
|
interrupts = result.get("__interrupt__", [])
|
||||||
|
assert len(interrupts) == 1
|
||||||
|
assert node_counter == 1
|
||||||
|
|
||||||
|
# invoke with None. this should NOT execute the node and the history should
|
||||||
|
# look the same as the first run
|
||||||
|
partial = graph.invoke(None, config=config)
|
||||||
|
second_history = list(graph.get_state_history(config))
|
||||||
|
remaining_interrupts = partial.get("__interrupt__", [])
|
||||||
|
assert len(remaining_interrupts) == 1
|
||||||
|
assert remaining_interrupts[0].value == "first"
|
||||||
|
assert node_counter == 1
|
||||||
|
|
||||||
|
# history should look the same for tasks and interrupts
|
||||||
|
print("first_history[0].interrupts: ", first_history[0].interrupts)
|
||||||
|
print("second_history[0].interrupts: ", second_history[0].interrupts)
|
||||||
|
print("first_history[0].tasks: ", first_history[0].tasks)
|
||||||
|
print("second_history[0].tasks: ", second_history[0].tasks)
|
||||||
|
assert first_history[0].interrupts == second_history[0].interrupts
|
||||||
|
assert first_history[0].tasks == second_history[0].tasks
|
||||||
|
|
||||||
|
# now resume the first interrupt with some value
|
||||||
|
partial = graph.invoke(Command(resume="weet"), config=config)
|
||||||
|
print("partial 3", partial)
|
||||||
|
third_history = list(graph.get_state_history(config))
|
||||||
|
remaining_interrupts = partial.get("__interrupt__", [])
|
||||||
|
assert len(remaining_interrupts) == 1
|
||||||
|
assert remaining_interrupts[0].value == "second"
|
||||||
|
assert node_counter == 2
|
||||||
|
|
||||||
|
# invoke with None again. the history should look the same as
|
||||||
|
# the third run
|
||||||
|
partial = graph.invoke(None, config=config)
|
||||||
|
print("partial 4", partial)
|
||||||
|
fourth_history = list(graph.get_state_history(config))
|
||||||
|
remaining_interrupts = partial.get("__interrupt__", [])
|
||||||
|
assert len(remaining_interrupts) == 1
|
||||||
|
assert node_counter == 2
|
||||||
|
|
||||||
|
print("\nthird_history[0].interrupts: ", third_history[0].interrupts)
|
||||||
|
print("fourth_history[0].interrupts: ", fourth_history[0].interrupts)
|
||||||
|
print("third_history[0].tasks: ", third_history[0].tasks)
|
||||||
|
print("fourth_history[0].tasks: ", fourth_history[0].tasks)
|
||||||
|
assert third_history[0].interrupts == fourth_history[0].interrupts
|
||||||
|
assert third_history[0].tasks == fourth_history[0].tasks
|
||||||
|
|
||||||
|
# resume the graph once more with a real value
|
||||||
|
partial = graph.invoke(Command(resume="bix"), config=config)
|
||||||
|
remaining_interrupts = partial.get("__interrupt__", [])
|
||||||
|
assert len(remaining_interrupts) == 0
|
||||||
|
assert node_counter == 3
|
||||||
|
|||||||
@@ -4726,7 +4726,7 @@ def test_send_dedupe_on_resume(
|
|||||||
assert len(history) == (4 if durability != "exit" else 1)
|
assert len(history) == (4 if durability != "exit" else 1)
|
||||||
|
|
||||||
# resume execution
|
# resume execution
|
||||||
assert graph.invoke(None, thread1, durability=durability) == [
|
assert graph.invoke(Command(resume=""), thread1, durability=durability) == [
|
||||||
"0",
|
"0",
|
||||||
"1",
|
"1",
|
||||||
"3.1",
|
"3.1",
|
||||||
|
|||||||
@@ -2545,7 +2545,7 @@ async def test_send_dedupe_on_resume(
|
|||||||
assert builder.nodes["2"].runnable.func.ticks == 3
|
assert builder.nodes["2"].runnable.func.ticks == 3
|
||||||
assert builder.nodes["flaky"].runnable.func.ticks == 1
|
assert builder.nodes["flaky"].runnable.func.ticks == 1
|
||||||
# resume execution
|
# resume execution
|
||||||
assert await graph.ainvoke(None, thread1, durability=durability) == [
|
assert await graph.ainvoke(Command(resume=""), thread1, durability=durability) == [
|
||||||
"0",
|
"0",
|
||||||
"1",
|
"1",
|
||||||
"3.1",
|
"3.1",
|
||||||
|
|||||||
Reference in New Issue
Block a user