Compare commits

...
6 changed files with 353 additions and 5 deletions
+7
View File
@@ -76,6 +76,7 @@ from langgraph.types import (
CachePolicy, CachePolicy,
Checkpointer, Checkpointer,
Command, Command,
OnInterruptHook,
RetryPolicy, RetryPolicy,
Send, Send,
ensure_valid_checkpointer, ensure_valid_checkpointer,
@@ -831,6 +832,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
interrupt_after: All | list[str] | None = None, interrupt_after: All | list[str] | None = None,
debug: bool = False, debug: bool = False,
name: str | None = None, name: str | None = None,
on_interrupt: OnInterruptHook | None = None,
) -> CompiledStateGraph[StateT, ContextT, InputT, OutputT]: ) -> CompiledStateGraph[StateT, ContextT, InputT, OutputT]:
"""Compiles the `StateGraph` into a `CompiledStateGraph` object. """Compiles the `StateGraph` into a `CompiledStateGraph` object.
@@ -850,6 +852,10 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
interrupt_after: An optional list of node names to interrupt after. interrupt_after: An optional list of node names to interrupt after.
debug: A flag indicating whether to enable debug mode. debug: A flag indicating whether to enable debug mode.
name: The name to use for the compiled graph. name: The name to use for the compiled graph.
on_interrupt: An optional callback that is invoked whenever the graph
execution is interrupted. Called with the list of `Interrupt` objects.
May be a sync function or an async coroutine function.
Returns: Returns:
CompiledStateGraph: The compiled `StateGraph`. CompiledStateGraph: The compiled `StateGraph`.
@@ -910,6 +916,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
store=store, store=store,
cache=cache, cache=cache,
name=name or "LangGraph", name=name or "LangGraph",
on_interrupt=on_interrupt,
) )
compiled.attach_node(START, None) compiled.attach_node(START, None)
+59 -4
View File
@@ -3,6 +3,7 @@ from __future__ import annotations
import asyncio import asyncio
import binascii import binascii
import concurrent.futures import concurrent.futures
import warnings
from collections import defaultdict, deque from collections import defaultdict, deque
from collections.abc import Callable, Iterator, Mapping, Sequence from collections.abc import Callable, Iterator, Mapping, Sequence
from contextlib import ( from contextlib import (
@@ -12,7 +13,7 @@ from contextlib import (
ExitStack, ExitStack,
) )
from datetime import datetime, timezone from datetime import datetime, timezone
from inspect import signature from inspect import iscoroutinefunction, signature
from types import TracebackType from types import TracebackType
from typing import ( from typing import (
Any, Any,
@@ -115,6 +116,8 @@ from langgraph.types import (
CachePolicy, CachePolicy,
Command, Command,
Durability, Durability,
Interrupt,
OnInterruptHook,
PregelExecutableTask, PregelExecutableTask,
RetryPolicy, RetryPolicy,
Send, Send,
@@ -157,6 +160,7 @@ class PregelLoop:
manager: None | AsyncParentRunManager | ParentRunManager manager: None | AsyncParentRunManager | ParentRunManager
interrupt_after: All | Sequence[str] interrupt_after: All | Sequence[str]
interrupt_before: All | Sequence[str] interrupt_before: All | Sequence[str]
on_interrupt: OnInterruptHook | None
durability: Durability durability: Durability
retry_policy: Sequence[RetryPolicy] retry_policy: Sequence[RetryPolicy]
cache_policy: CachePolicy | None cache_policy: CachePolicy | None
@@ -226,6 +230,7 @@ class PregelLoop:
migrate_checkpoint: Callable[[Checkpoint], None] | None = None, migrate_checkpoint: Callable[[Checkpoint], None] | None = None,
retry_policy: Sequence[RetryPolicy] = (), retry_policy: Sequence[RetryPolicy] = (),
cache_policy: CachePolicy | None = None, cache_policy: CachePolicy | None = None,
on_interrupt: OnInterruptHook | None = None,
) -> None: ) -> None:
self.stream = stream self.stream = stream
self.config = config self.config = config
@@ -242,6 +247,7 @@ class PregelLoop:
self.stream_keys = stream_keys self.stream_keys = stream_keys
self.interrupt_after = interrupt_after self.interrupt_after = interrupt_after
self.interrupt_before = interrupt_before self.interrupt_before = interrupt_before
self.on_interrupt = on_interrupt
self.manager = manager self.manager = manager
self.is_nested = CONFIG_KEY_TASK_ID in self.config.get(CONF, {}) self.is_nested = CONFIG_KEY_TASK_ID in self.config.get(CONF, {})
self.skip_done_tasks = CONFIG_KEY_CHECKPOINT_ID not in config[CONF] self.skip_done_tasks = CONFIG_KEY_CHECKPOINT_ID not in config[CONF]
@@ -865,14 +871,36 @@ class PregelLoop:
[{INTERRUPT: cast(GraphInterrupt, exc_value).args[0]}] [{INTERRUPT: cast(GraphInterrupt, exc_value).args[0]}]
), ),
) )
# save final output # save final output first, so graph state is consistent even
# if the on_interrupt hook raises
self.output = read_channels(self.channels, self.output_keys) self.output = read_channels(self.channels, self.output_keys)
# call on_interrupt hook
if self.on_interrupt is not None:
interrupts: list[Interrupt] = (
list(cast(GraphInterrupt, exc_value).args[0])
if exc_value is not None and exc_value.args and exc_value.args[0]
else []
)
self._call_on_interrupt(interrupts)
# suppress interrupt # suppress interrupt
return True return True
elif exc_type is None: elif exc_type is None:
# save final output # save final output
self.output = read_channels(self.channels, self.output_keys) self.output = read_channels(self.channels, self.output_keys)
def _call_on_interrupt(self, interrupts: list[Interrupt]) -> None:
"""Call the on_interrupt hook synchronously."""
if self.on_interrupt is None:
return
if iscoroutinefunction(self.on_interrupt):
warnings.warn(
"Async on_interrupt hook cannot be called from sync graph execution. "
"Use a sync function or run the graph with astream/ainvoke.",
stacklevel=2,
)
return
self.on_interrupt(interrupts)
def _emit( def _emit(
self, self,
mode: StreamMode, mode: StreamMode,
@@ -985,6 +1013,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
migrate_checkpoint: Callable[[Checkpoint], None] | None = None, migrate_checkpoint: Callable[[Checkpoint], None] | None = None,
retry_policy: Sequence[RetryPolicy] = (), retry_policy: Sequence[RetryPolicy] = (),
cache_policy: CachePolicy | None = None, cache_policy: CachePolicy | None = None,
on_interrupt: OnInterruptHook | None = None,
) -> None: ) -> None:
super().__init__( super().__init__(
input, input,
@@ -1006,6 +1035,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
retry_policy=retry_policy, retry_policy=retry_policy,
cache_policy=cache_policy, cache_policy=cache_policy,
durability=durability, durability=durability,
on_interrupt=on_interrupt,
) )
self.stack = ExitStack() self.stack = ExitStack()
if checkpointer: if checkpointer:
@@ -1161,6 +1191,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
migrate_checkpoint: Callable[[Checkpoint], None] | None = None, migrate_checkpoint: Callable[[Checkpoint], None] | None = None,
retry_policy: Sequence[RetryPolicy] = (), retry_policy: Sequence[RetryPolicy] = (),
cache_policy: CachePolicy | None = None, cache_policy: CachePolicy | None = None,
on_interrupt: OnInterruptHook | None = None,
) -> None: ) -> None:
super().__init__( super().__init__(
input, input,
@@ -1182,6 +1213,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
retry_policy=retry_policy, retry_policy=retry_policy,
cache_policy=cache_policy, cache_policy=cache_policy,
durability=durability, durability=durability,
on_interrupt=on_interrupt,
) )
self.stack = AsyncExitStack() self.stack = AsyncExitStack()
if checkpointer: if checkpointer:
@@ -1257,6 +1289,18 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
}, },
) )
_deferred_on_interrupt_args: list[Interrupt] | None = None
def _call_on_interrupt(self, interrupts: list[Interrupt]) -> None:
"""Override for async loop: defer async hooks to __aexit__."""
if self.on_interrupt is None:
return
if iscoroutinefunction(self.on_interrupt):
# Defer async hooks — they will be awaited in __aexit__
self._deferred_on_interrupt_args = interrupts
else:
self.on_interrupt(interrupts)
# context manager # context manager
async def __aenter__(self) -> Self: async def __aenter__(self) -> Self:
@@ -1315,14 +1359,25 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
exc_value: BaseException | None, exc_value: BaseException | None,
traceback: TracebackType | None, traceback: TracebackType | None,
) -> bool | None: ) -> bool | None:
# unwind stack # unwind stack (calls _suppress_interrupt synchronously)
exit_task = asyncio.create_task( exit_task = asyncio.create_task(
self.stack.__aexit__(exc_type, exc_value, traceback) self.stack.__aexit__(exc_type, exc_value, traceback)
) )
try: try:
return await exit_task result = await exit_task
except asyncio.CancelledError as e: except asyncio.CancelledError as e:
# Bubble up the exit task upon cancellation to permit the API # Bubble up the exit task upon cancellation to permit the API
# consumer to await it before e.g., reusing the DB connection. # consumer to await it before e.g., reusing the DB connection.
e.args = (*e.args, exit_task) e.args = (*e.args, exit_task)
raise raise
# Await deferred async on_interrupt hook (set by _call_on_interrupt)
if (
self._deferred_on_interrupt_args is not None
and self.on_interrupt is not None
):
interrupts = self._deferred_on_interrupt_args
self._deferred_on_interrupt_args = None
coro = self.on_interrupt(interrupts)
if coro is not None:
await coro
return result
+11
View File
@@ -138,6 +138,7 @@ from langgraph.types import (
Command, Command,
Durability, Durability,
Interrupt, Interrupt,
OnInterruptHook,
Send, Send,
StateSnapshot, StateSnapshot,
StateUpdate, StateUpdate,
@@ -622,6 +623,12 @@ class Pregel(
context_schema: type[ContextT] | None = None context_schema: type[ContextT] | None = None
"""Specifies the schema for the context object that will be passed to the workflow.""" """Specifies the schema for the context object that will be passed to the workflow."""
on_interrupt: OnInterruptHook | None = None
"""Optional callback invoked when the graph execution is interrupted.
Called with the list of `Interrupt` objects whenever the graph pauses.
May be a sync or async callable."""
config: RunnableConfig | None = None config: RunnableConfig | None = None
name: str = "LangGraph" name: str = "LangGraph"
@@ -652,6 +659,7 @@ class Pregel(
config: RunnableConfig | None = None, config: RunnableConfig | None = None,
trigger_to_nodes: Mapping[str, Sequence[str]] | None = None, trigger_to_nodes: Mapping[str, Sequence[str]] | None = None,
name: str = "LangGraph", name: str = "LangGraph",
on_interrupt: OnInterruptHook | None = None,
**deprecated_kwargs: Unpack[DeprecatedKwargs], **deprecated_kwargs: Unpack[DeprecatedKwargs],
) -> None: ) -> None:
if ( if (
@@ -695,6 +703,7 @@ class Pregel(
) )
self.cache_policy = cache_policy self.cache_policy = cache_policy
self.context_schema = context_schema self.context_schema = context_schema
self.on_interrupt = on_interrupt
self.config = config self.config = config
self.trigger_to_nodes = trigger_to_nodes or {} self.trigger_to_nodes = trigger_to_nodes or {}
self.name = name self.name = name
@@ -2599,6 +2608,7 @@ class Pregel(
migrate_checkpoint=self._migrate_checkpoint, migrate_checkpoint=self._migrate_checkpoint,
retry_policy=self.retry_policy, retry_policy=self.retry_policy,
cache_policy=self.cache_policy, cache_policy=self.cache_policy,
on_interrupt=self.on_interrupt,
) as loop: ) as loop:
# create runner # create runner
runner = PregelRunner( runner = PregelRunner(
@@ -2908,6 +2918,7 @@ class Pregel(
migrate_checkpoint=self._migrate_checkpoint, migrate_checkpoint=self._migrate_checkpoint,
retry_policy=self.retry_policy, retry_policy=self.retry_policy,
cache_policy=self.cache_policy, cache_policy=self.cache_policy,
on_interrupt=self.on_interrupt,
) as loop: ) as loop:
# create runner # create runner
runner = PregelRunner( runner = PregelRunner(
+14 -1
View File
@@ -2,7 +2,7 @@ from __future__ import annotations
import sys import sys
from collections import deque from collections import deque
from collections.abc import Callable, Hashable, Sequence from collections.abc import Awaitable, Callable, Hashable, Sequence
from dataclasses import asdict, dataclass from dataclasses import asdict, dataclass
from typing import ( from typing import (
TYPE_CHECKING, TYPE_CHECKING,
@@ -56,6 +56,7 @@ __all__ = (
"Durability", "Durability",
"interrupt", "interrupt",
"Overwrite", "Overwrite",
"OnInterruptHook",
"ensure_valid_checkpointer", "ensure_valid_checkpointer",
) )
@@ -109,6 +110,18 @@ StreamWriter = Callable[[Any], None]
Always injected into nodes if requested as a keyword argument, but it's a no-op Always injected into nodes if requested as a keyword argument, but it's a no-op
when not using `stream_mode="custom"`.""" when not using `stream_mode="custom"`."""
OnInterruptHook = (
Callable[[list["Interrupt"]], None] | Callable[[list["Interrupt"]], Awaitable[None]]
)
"""Callback invoked when a graph execution is interrupted.
Called with the list of `Interrupt` objects whenever the graph pauses due to
an `interrupt()` call or `interrupt_before`/`interrupt_after` configuration.
May be a regular function or an async coroutine function. Async hooks are
awaited in async graph execution; in sync execution only sync hooks are called.
"""
_DC_KWARGS = {"kw_only": True, "slots": True, "frozen": True} _DC_KWARGS = {"kw_only": True, "slots": True, "frozen": True}
+148
View File
@@ -8893,3 +8893,151 @@ def test_fork_does_not_apply_pending_writes(
# Should be: 1 (input) + 20 (forked node_a) + 100 (node_b) = 121 # Should be: 1 (input) + 20 (forked node_a) + 100 (node_b) = 121
assert result == {"value": 121} assert result == {"value": 121}
def test_on_interrupt_hook_with_interrupt_call(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
"""Test that on_interrupt hook fires when interrupt() is called in a node."""
hook_calls: list[list[Interrupt]] = []
def my_on_interrupt(interrupts: list[Interrupt]) -> None:
hook_calls.append(interrupts)
class State(TypedDict):
value: str
def ask_human(state: State) -> dict:
answer = interrupt("what should I do?")
return {"value": answer}
builder = StateGraph(State)
builder.add_node("ask", ask_human)
builder.add_edge(START, "ask")
graph = builder.compile(
checkpointer=sync_checkpointer,
on_interrupt=my_on_interrupt,
)
config = {"configurable": {"thread_id": "1"}}
# First invocation: should trigger interrupt and call the hook
result = list(graph.stream({"value": ""}, config))
assert len(result) == 1
assert "__interrupt__" in result[0]
# Hook should have been called once with the interrupt data
assert len(hook_calls) == 1
assert len(hook_calls[0]) == 1
assert hook_calls[0][0].value == "what should I do?"
# Resume — no new interrupt, hook should not fire again
hook_calls.clear()
result = list(graph.stream(Command(resume="do this"), config))
assert any("ask" in chunk for chunk in result)
assert len(hook_calls) == 0
def test_on_interrupt_hook_with_interrupt_before(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
"""Test that on_interrupt hook fires for interrupt_before config."""
hook_calls: list[list[Interrupt]] = []
def my_on_interrupt(interrupts: list[Interrupt]) -> None:
hook_calls.append(interrupts)
class State(TypedDict):
value: int
def add_one(state: State) -> dict:
return {"value": state["value"] + 1}
builder = StateGraph(State)
builder.add_node("add_one", add_one)
builder.add_edge(START, "add_one")
graph = builder.compile(
checkpointer=sync_checkpointer,
interrupt_before=["add_one"],
on_interrupt=my_on_interrupt,
)
config = {"configurable": {"thread_id": "1"}}
# Should interrupt before add_one runs
result = list(graph.stream({"value": 0}, config))
assert any("__interrupt__" in chunk for chunk in result)
# Hook should have been called (empty interrupt list for config-level interrupts)
assert len(hook_calls) == 1
assert hook_calls[0] == []
def test_on_interrupt_hook_not_called_without_interrupt(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
"""Test that on_interrupt hook is NOT called when graph completes normally."""
hook_calls: list[list[Interrupt]] = []
def my_on_interrupt(interrupts: list[Interrupt]) -> None:
hook_calls.append(interrupts)
class State(TypedDict):
value: int
def add_one(state: State) -> dict:
return {"value": state["value"] + 1}
builder = StateGraph(State)
builder.add_node("add_one", add_one)
builder.add_edge(START, "add_one")
graph = builder.compile(
checkpointer=sync_checkpointer,
on_interrupt=my_on_interrupt,
)
config = {"configurable": {"thread_id": "1"}}
result = graph.invoke({"value": 0}, config)
assert result == {"value": 1}
# Hook should NOT have been called
assert len(hook_calls) == 0
def test_on_interrupt_hook_exception_propagates(
sync_checkpointer: BaseCheckpointSaver,
) -> None:
"""Test that exceptions in the on_interrupt hook propagate to the caller."""
def bad_hook(interrupts: list[Interrupt]) -> None:
raise RuntimeError("hook exploded")
class State(TypedDict):
value: str
def ask(state: State) -> dict:
answer = interrupt("question")
return {"value": answer}
builder = StateGraph(State)
builder.add_node("ask", ask)
builder.add_edge(START, "ask")
graph = builder.compile(
checkpointer=sync_checkpointer,
on_interrupt=bad_hook,
)
config = {"configurable": {"thread_id": "1"}}
# Hook error should propagate
with pytest.raises(RuntimeError, match="hook exploded"):
list(graph.stream({"value": ""}, config))
# Graph state should still be checkpointed and resumable despite the hook error
result = list(graph.stream(Command(resume="answer"), config))
assert any("ask" in chunk for chunk in result)
+114
View File
@@ -9345,3 +9345,117 @@ async def test_fork_does_not_apply_pending_writes(
# 1 (input) + 20 (forked node_a) + 100 (node_b) = 121 # 1 (input) + 20 (forked node_a) + 100 (node_b) = 121
assert result == {"value": 121} assert result == {"value": 121}
async def test_on_interrupt_hook_async_with_interrupt_call(
async_checkpointer: BaseCheckpointSaver,
) -> None:
"""Test that an async on_interrupt hook fires when interrupt() is called."""
hook_calls: list[list[Interrupt]] = []
async def my_on_interrupt(interrupts: list[Interrupt]) -> None:
hook_calls.append(interrupts)
class State(TypedDict):
value: str
def ask_human(state: State) -> dict:
answer = interrupt("what should I do?")
return {"value": answer}
builder = StateGraph(State)
builder.add_node("ask", ask_human)
builder.add_edge(START, "ask")
graph = builder.compile(
checkpointer=async_checkpointer,
on_interrupt=my_on_interrupt,
)
config = {"configurable": {"thread_id": "1"}}
# First invocation: should trigger interrupt and call the async hook
result = [chunk async for chunk in graph.astream({"value": ""}, config)]
assert len(result) == 1
assert "__interrupt__" in result[0]
# Hook should have been called once with the interrupt data
assert len(hook_calls) == 1
assert len(hook_calls[0]) == 1
assert hook_calls[0][0].value == "what should I do?"
# Resume — no new interrupt, hook should not fire again
hook_calls.clear()
result = [chunk async for chunk in graph.astream(Command(resume="do this"), config)]
assert any("ask" in chunk for chunk in result)
assert len(hook_calls) == 0
async def test_on_interrupt_hook_sync_in_async_graph(
async_checkpointer: BaseCheckpointSaver,
) -> None:
"""Test that a sync on_interrupt hook works in async graph execution."""
hook_calls: list[list[Interrupt]] = []
def my_sync_hook(interrupts: list[Interrupt]) -> None:
hook_calls.append(interrupts)
class State(TypedDict):
value: str
def ask_human(state: State) -> dict:
answer = interrupt("question?")
return {"value": answer}
builder = StateGraph(State)
builder.add_node("ask", ask_human)
builder.add_edge(START, "ask")
graph = builder.compile(
checkpointer=async_checkpointer,
on_interrupt=my_sync_hook,
)
config = {"configurable": {"thread_id": "1"}}
result = [chunk async for chunk in graph.astream({"value": ""}, config)]
assert "__interrupt__" in result[0]
# Sync hook should work fine in async execution
assert len(hook_calls) == 1
assert hook_calls[0][0].value == "question?"
async def test_on_interrupt_hook_async_exception_propagates(
async_checkpointer: BaseCheckpointSaver,
) -> None:
"""Test that exceptions in the async on_interrupt hook propagate."""
async def bad_hook(interrupts: list[Interrupt]) -> None:
raise RuntimeError("async hook exploded")
class State(TypedDict):
value: str
def ask(state: State) -> dict:
answer = interrupt("question")
return {"value": answer}
builder = StateGraph(State)
builder.add_node("ask", ask)
builder.add_edge(START, "ask")
graph = builder.compile(
checkpointer=async_checkpointer,
on_interrupt=bad_hook,
)
config = {"configurable": {"thread_id": "1"}}
# Hook error should propagate
with pytest.raises(RuntimeError, match="async hook exploded"):
[chunk async for chunk in graph.astream({"value": ""}, config)]
# Graph state should still be checkpointed and resumable despite the hook error
result = [chunk async for chunk in graph.astream(Command(resume="answer"), config)]
assert any("ask" in chunk for chunk in result)