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

Enable raising Interrupt from inside a node, add list of current interrupts to get_state (#1354)

Browse Source 
* Enable raising Interrupt from inside a node, add list of current interrupts to get_state

* Lint

* Allow multiple interrupt values in exception

* Better typings

* Fix some tests

* Fix up

* Lint

* Add test

* Lint

* WIP stay=True

* Fix tests for get_state

* Remove ids

* Fix step count

* 999

* less fun

* Undo

* Update debug interface

* Remove ability to pass multiple values

* Undo
This commit is contained in:
Nuno Campos
2024-08-21 16:19:15 +00:00
committed by GitHub
parent ebf060675d
commit f8b053a286
10 changed files with 473 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/constants.py
+10 -1
View File
@@ -1,10 +1,12 @@
from typing import Any
from dataclasses import dataclass
from typing import Any, Literal
INPUT = "__input__"
CONFIG_KEY_SEND = "__pregel_send"
CONFIG_KEY_READ = "__pregel_read"
CONFIG_KEY_CHECKPOINTER = "__pregel_checkpointer"
CONFIG_KEY_RESUMING = "__pregel_resuming"
CONFIG_KEY_TASK_ID = "__pregel_task_id"
INTERRUPT = "__interrupt__"
ERROR = "__error__"
TASKS = "__pregel_tasks"
@@ -16,6 +18,7 @@ RESERVED = {
CONFIG_KEY_READ,
CONFIG_KEY_CHECKPOINTER,
CONFIG_KEY_RESUMING,
CONFIG_KEY_TASK_ID,
INPUT,
}
TAG_HIDDEN = "langsmith:hidden"
@@ -93,3 +96,9 @@ class Send:
and self.node == value.node
and self.arg == value.arg
)
@dataclass
class Interrupt:
when: Literal["before", "during", "after"]
value: Any = None
libs/langgraph/langgraph/errors.py
+13 -1
View File
@@ -1,4 +1,7 @@
from typing import Any
from langgraph.checkpoint.base import EmptyChannelError
from langgraph.constants import Interrupt
class GraphRecursionError(RecursionError):
@@ -29,7 +32,15 @@ class InvalidUpdateError(Exception):
class GraphInterrupt(Exception):
"""Raised when a subgraph is interrupted."""
pass
def __init__(self, interrupts: list[Interrupt]) -> None:
super().__init__(interrupts)
class NodeInterrupt(GraphInterrupt):
"""Raised by a node to interrupt execution."""
def __init__(self, value: Any) -> None:
super().__init__([Interrupt("during", value)])
class EmptyInputError(Exception):
@@ -42,6 +53,7 @@ __all__ = [
"GraphRecursionError",
"InvalidUpdateError",
"GraphInterrupt",
"NodeInterrupt",
"EmptyInputError",
"EmptyChannelError",
]
libs/langgraph/langgraph/pregel/__init__.py
+109 -29
View File
@@ -70,8 +70,9 @@ from langgraph.constants import (
CONFIG_KEY_SEND,
ERROR,
INTERRUPT,
Interrupt,
)
from langgraph.errors import GraphRecursionError, InvalidUpdateError
from langgraph.errors import GraphInterrupt, GraphRecursionError, InvalidUpdateError
from langgraph.managed.base import (
AsyncManagedValuesManager,
ManagedValuesManager,
@@ -82,6 +83,7 @@ from langgraph.pregel.algo import (
apply_writes,
local_read,
prepare_next_tasks,
should_interrupt,
)
from langgraph.pregel.debug import (
map_debug_task_results,
@@ -381,6 +383,7 @@ class Pregel(
saved.metadata.get("step", -1) + 1 if saved else -1,
for_execution=False,
)
return StateSnapshot(
read_channels(channels, self.stream_channels_asis),
tuple(t.name for t in next_tasks),
@@ -550,7 +553,7 @@ class Pregel(
saved = self.checkpointer.get_tuple(config)
checkpoint = copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
checkpoint_previous_versions = (
saved.checkpoint["channel_versions"] if saved else {}
saved.checkpoint["channel_versions"].copy() if saved else {}
)
step = saved.metadata.get("step", -1) if saved else -1
# merge configurable fields with previous checkpoint config
@@ -576,7 +579,7 @@ class Pregel(
create_checkpoint(checkpoint, None, step),
{
"source": "update",
"step": step,
"step": step + 1,
"writes": {},
},
{},
@@ -606,7 +609,11 @@ class Pregel(
if as_node not in self.nodes:
raise InvalidUpdateError(f"Node {as_node} does not exist")
# update channels
with ChannelsManager(self.channels, checkpoint, config) as channels:
with ChannelsManager(
self.channels, checkpoint, config
) as channels, ManagedValuesManager(
self.managed_values_dict, ensure_config(config)
) as managed:
# create task to run all writers of the chosen node
writers = self.nodes[as_node].get_writers()
if not writers:
@@ -640,19 +647,43 @@ class Pregel(
apply_writes(
checkpoint, channels, [task], self.checkpointer.get_next_version
)
new_versions = get_new_channel_versions(
checkpoint_previous_versions, checkpoint["channel_versions"]
)
checkpoint = create_checkpoint(checkpoint, channels, step + 1)
# check interrupt before
if tasks := should_interrupt(
checkpoint,
self.interrupt_before_nodes,
prepare_next_tasks(
checkpoint,
self.nodes,
channels,
managed,
config,
step + 2,
for_execution=False,
),
):
for t in tasks:
self.checkpointer.put_writes(
{
"configurable": {
**checkpoint_config["configurable"],
"checkpoint_id": checkpoint["id"],
}
},
[(INTERRUPT, Interrupt("before"))],
t.id,
)
return self.checkpointer.put(
checkpoint_config,
create_checkpoint(checkpoint, channels, step + 1),
checkpoint,
{
"source": "update",
"step": step + 1,
"writes": {as_node: values},
},
new_versions,
get_new_channel_versions(
checkpoint_previous_versions, checkpoint["channel_versions"]
),
)
async def aupdate_state(
@@ -668,7 +699,7 @@ class Pregel(
saved = await self.checkpointer.aget_tuple(config)
checkpoint = copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
checkpoint_previous_versions = (
saved.checkpoint["channel_versions"] if saved else {}
saved.checkpoint["channel_versions"].copy() if saved else {}
)
step = saved.metadata.get("step", -1) if saved else -1
# merge configurable fields with previous checkpoint config
@@ -694,7 +725,7 @@ class Pregel(
create_checkpoint(checkpoint, None, step),
{
"source": "update",
"step": step,
"step": step + 1,
"writes": {},
},
{},
@@ -722,7 +753,11 @@ class Pregel(
if as_node not in self.nodes:
raise InvalidUpdateError(f"Node {as_node} does not exist")
# update channels, acting as the chosen node
async with AsyncChannelsManager(self.channels, checkpoint, config) as channels:
async with AsyncChannelsManager(
self.channels, checkpoint, config
) as channels, AsyncManagedValuesManager(
self.managed_values_dict, ensure_config(config)
) as managed:
# create task to run all writers of the chosen node
writers = self.nodes[as_node].get_writers()
if not writers:
@@ -756,19 +791,47 @@ class Pregel(
apply_writes(
checkpoint, channels, [task], self.checkpointer.get_next_version
)
new_versions = get_new_channel_versions(
checkpoint_previous_versions, checkpoint["channel_versions"]
)
checkpoint = create_checkpoint(checkpoint, channels, step + 1)
# check interrupt before
if tasks := should_interrupt(
checkpoint,
self.interrupt_before_nodes,
prepare_next_tasks(
checkpoint,
self.nodes,
channels,
managed,
config,
step + 2,
for_execution=False,
),
):
await asyncio.gather(
*(
self.checkpointer.aput_writes(
{
"configurable": {
**checkpoint_config["configurable"],
"checkpoint_id": checkpoint["id"],
}
},
[(INTERRUPT, Interrupt("before"))],
t.id,
)
for t in tasks
)
)
return await self.checkpointer.aput(
checkpoint_config,
create_checkpoint(checkpoint, channels, step + 1),
checkpoint,
{
"source": "update",
"step": step + 1,
"writes": {as_node: values},
},
new_versions,
get_new_channel_versions(
checkpoint_previous_versions, checkpoint["channel_versions"]
),
)
def _defaults(
@@ -975,6 +1038,7 @@ class Pregel(
for task in loop.tasks
if not task.writes
}
all_futures = futures.copy()
end_time = (
self.step_timeout + time.monotonic()
if self.step_timeout
@@ -998,7 +1062,12 @@ class Pregel(
task = futures.pop(fut)
if exc := _exception(fut):
# save error to checkpointer
loop.put_writes(task.id, [(ERROR, exc)])
if isinstance(exc, GraphInterrupt):
loop.put_writes(
task.id, [(INTERRUPT, i) for i in exc.args[0]]
)
else:
loop.put_writes(task.id, [(ERROR, exc)])
else:
# save task writes to checkpointer
loop.put_writes(task.id, task.writes)
@@ -1026,7 +1095,7 @@ class Pregel(
break
# panic on failure or timeout
_panic_or_proceed(done, inflight, loop.step)
_panic_or_proceed(all_futures, loop.step)
# don't keep futures around in memory longer than needed
del done, inflight, futures
# debug flag
@@ -1223,6 +1292,7 @@ class Pregel(
for task in loop.tasks
if not task.writes
}
all_futures = futures.copy()
end_time = (
self.step_timeout + aioloop.time()
if self.step_timeout
@@ -1244,7 +1314,12 @@ class Pregel(
task = futures.pop(fut)
if exc := _exception(fut):
# save error to checkpointer
loop.put_writes(task.id, [(ERROR, exc)])
if isinstance(exc, GraphInterrupt):
loop.put_writes(
task.id, [(INTERRUPT, i) for i in exc.args[0]]
)
else:
loop.put_writes(task.id, [(ERROR, exc)])
else:
# save task writes to checkpointer
loop.put_writes(task.id, task.writes)
@@ -1274,7 +1349,7 @@ class Pregel(
break
# panic on failure or timeout
_panic_or_proceed(done, inflight, loop.step, asyncio.TimeoutError)
_panic_or_proceed(all_futures, loop.step, asyncio.TimeoutError)
# don't keep futures around in memory longer than needed
del done, inflight, futures
# debug flag
@@ -1421,9 +1496,8 @@ def _should_stop_others(
for fut in done:
if fut.cancelled():
return True
if fut.exception() is not None:
# TODO don't stop others if exception is interrupt
return True
if exc := fut.exception():
return not isinstance(exc, GraphInterrupt)
else:
return False
@@ -1441,14 +1515,20 @@ def _exception(
def _panic_or_proceed(
done: Union[set[concurrent.futures.Future[Any]], set[asyncio.Task[Any]]],
inflight: Union[set[concurrent.futures.Future[Any]], set[asyncio.Task[Any]]],
futs: Union[set[concurrent.futures.Future[Any]], set[asyncio.Task[Any]]],
step: int,
timeout_exc_cls: Type[Exception] = TimeoutError,
) -> None:
done: set[Union[concurrent.futures.Future[Any], asyncio.Task[Any]]] = set()
inflight: set[Union[concurrent.futures.Future[Any], asyncio.Task[Any]]] = set()
for fut in futs:
if fut.done():
done.add(fut)
else:
inflight.add(fut)
while done:
# if any task failed
if exc := done.pop().exception():
if exc := _exception(done.pop()):
# cancel all pending tasks
while inflight:
inflight.pop().cancel()
libs/langgraph/langgraph/pregel/algo.py
+15 -10
View File
@@ -38,6 +38,7 @@ from langgraph.constants import (
CONFIG_KEY_READ,
CONFIG_KEY_RESUMING,
CONFIG_KEY_SEND,
CONFIG_KEY_TASK_ID,
INTERRUPT,
RESERVED,
TAG_HIDDEN,
@@ -68,26 +69,28 @@ def should_interrupt(
checkpoint: Checkpoint,
interrupt_nodes: Union[All, Sequence[str]],
tasks: list[PregelExecutableTask],
) -> bool:
) -> list[PregelExecutableTask]:
version_type = type(next(iter(checkpoint["channel_versions"].values()), None))
null_version = version_type()
seen = checkpoint["versions_seen"].get(INTERRUPT, {})
# interrupt if any channel has been updated since last interrupt
any_updates_since_prev_interrupt = any(
version > seen.get(chan, null_version)
for chan, version in checkpoint["channel_versions"].items()
)
# and any triggered node is in interrupt_nodes list
return (
# interrupt if any channel has been updated since last interrupt
any(
version > seen.get(chan, null_version)
for chan, version in checkpoint["channel_versions"].items()
)
# and any triggered node is in interrupt_nodes list
and any(
task.name
[
task
for task in tasks
if (
(not task.config or TAG_HIDDEN not in task.config.get("tags"))
if interrupt_nodes == "*"
else task.name in interrupt_nodes
)
)
]
if any_updates_since_prev_interrupt
else []
)
@@ -308,6 +311,7 @@ def prepare_next_tasks(
else None
),
configurable={
CONFIG_KEY_TASK_ID: task_id,
# deque.extend is thread-safe
CONFIG_KEY_SEND: partial(
local_write, writes.extend, processes, channels
@@ -398,6 +402,7 @@ def prepare_next_tasks(
else None
),
configurable={
CONFIG_KEY_TASK_ID: task_id,
# deque.extend is thread-safe
CONFIG_KEY_SEND: partial(
local_write, writes.extend, processes, channels
libs/langgraph/langgraph/pregel/debug.py
+13 -5
View File
@@ -10,7 +10,7 @@ from langchain_core.utils.input import get_bolded_text, get_colored_text
from langgraph.channels.base import BaseChannel
from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, PendingWrite
from langgraph.constants import ERROR, TAG_HIDDEN
from langgraph.constants import ERROR, INTERRUPT, TAG_HIDDEN
from langgraph.pregel.io import read_channels
from langgraph.pregel.types import PregelExecutableTask, PregelTask
@@ -32,6 +32,7 @@ class CheckpointTask(TypedDict):
id: str
name: str
error: Optional[str]
interrupts: list[dict]
class CheckpointPayload(TypedDict):
@@ -149,6 +150,7 @@ def map_debug_checkpoint(
else {
"id": t.id,
"name": t.name,
"interrupts": t.interrupts,
}
for t in tasks_w_writes(tasks, pending_writes)
],
@@ -190,8 +192,11 @@ def print_step_writes(
def print_step_checkpoint(
step: int, channels: Mapping[str, BaseChannel], whitelist: Sequence[str]
metadata: CheckpointMetadata,
channels: Mapping[str, BaseChannel],
whitelist: Sequence[str],
) -> None:
step = metadata["step"]
print(
f"{get_colored_text(f'[{step}:checkpoint]', color='blue')} "
+ get_bolded_text(f"State at the end of step {step}:\n")
@@ -203,6 +208,7 @@ def tasks_w_writes(
tasks: list[PregelExecutableTask],
pending_writes: Optional[list[PendingWrite]],
) -> tuple[PregelTask, ...]:
pending_writes = pending_writes or []
return tuple(
PregelTask(
task.id,
@@ -210,12 +216,14 @@ def tasks_w_writes(
next(
(
exc
for tid, n, exc in pending_writes or []
if tid == task.id
if n == ERROR
for tid, n, exc in pending_writes
if tid == task.id and n == ERROR
),
None,
),
tuple(
v for tid, n, v in pending_writes if tid == task.id and n == INTERRUPT
),
)
for task in tasks
)
libs/langgraph/langgraph/pregel/loop.py
+15 -8
View File
@@ -45,6 +45,7 @@ from langgraph.constants import (
ERROR,
INPUT,
INTERRUPT,
Interrupt,
)
from langgraph.errors import EmptyInputError, GraphInterrupt
from langgraph.managed.base import (
@@ -206,10 +207,13 @@ class PregelLoop:
}
)
# after execution, check if we should interrupt
if should_interrupt(self.checkpoint, interrupt_after, self.tasks):
if tasks := should_interrupt(self.checkpoint, interrupt_after, self.tasks):
self.status = "interrupt_after"
interrupts = [(t.id, Interrupt("after")) for t in tasks]
for tid, interrupt in interrupts:
self.put_writes(tid, [(INTERRUPT, interrupt)])
if self.is_nested:
raise GraphInterrupt(self)
raise GraphInterrupt([i[1] for i in interrupts])
else:
return False
else:
@@ -258,7 +262,7 @@ class PregelLoop:
# if there are pending writes from a previous loop, apply them
if self.checkpoint_pending_writes:
for tid, k, v in self.checkpoint_pending_writes:
if k == ERROR: # TODO same for INTERRUPT
if k in (ERROR, INTERRUPT):
continue
if task := next((t for t in self.tasks if t.id == tid), None):
task.writes.append((k, v))
@@ -273,10 +277,13 @@ class PregelLoop:
)
# before execution, check if we should interrupt
if should_interrupt(self.checkpoint, interrupt_before, self.tasks):
if tasks := should_interrupt(self.checkpoint, interrupt_before, self.tasks):
self.status = "interrupt_before"
interrupts = [(t.id, Interrupt("before")) for t in tasks]
for tid, interrupt in interrupts:
self.put_writes(tid, [(INTERRUPT, interrupt)])
if self.is_nested:
raise GraphInterrupt()
raise GraphInterrupt([i[1] for i in interrupts])
else:
return False
@@ -394,7 +401,7 @@ class PregelLoop:
exc_value: Optional[BaseException],
traceback: Optional[TracebackType],
) -> Optional[bool]:
if exc_type is GraphInterrupt and not self.is_nested:
if isinstance(exc_value, GraphInterrupt) and not self.is_nested:
return True
@@ -409,7 +416,6 @@ class SyncPregelLoop(PregelLoop, ContextManager):
) -> None:
super().__init__(input, config=config, checkpointer=checkpointer, graph=graph)
self.stack = ExitStack()
self.stack.push(self._suppress_interrupt)
if checkpointer:
self.checkpointer_get_next_version = checkpointer.get_next_version
self.checkpointer_put_writes = checkpointer.put_writes
@@ -457,6 +463,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
self.managed = self.stack.enter_context(
ManagedValuesManager(self.graph.managed_values_dict, self.config)
)
self.stack.push(self._suppress_interrupt)
self.status = "pending"
self.step = self.checkpoint_metadata["step"] + 1
self.stop = self.step + self.config["recursion_limit"] + 1
@@ -486,7 +493,6 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
) -> None:
super().__init__(input, config=config, checkpointer=checkpointer, graph=graph)
self.stack = AsyncExitStack()
self.stack.push(self._suppress_interrupt)
if checkpointer:
self.checkpointer_get_next_version = checkpointer.get_next_version
self.checkpointer_put_writes = checkpointer.aput_writes
@@ -536,6 +542,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
self.managed = await self.stack.enter_async_context(
AsyncManagedValuesManager(self.graph.managed_values_dict, self.config)
)
self.stack.push(self._suppress_interrupt)
self.status = "pending"
self.step = self.checkpoint_metadata["step"] + 1
self.stop = self.step + self.config["recursion_limit"] + 1
libs/langgraph/langgraph/pregel/types.py
+2
View File
@@ -4,6 +4,7 @@ from typing import Any, Callable, Literal, NamedTuple, Optional, Type, Union
from langchain_core.runnables import Runnable, RunnableConfig
from langgraph.checkpoint.base import CheckpointMetadata
from langgraph.constants import Interrupt
def default_retry_on(exc: Exception) -> bool:
@@ -60,6 +61,7 @@ class PregelTask(NamedTuple):
id: str
name: str
error: Optional[Exception] = None
interrupts: tuple[Interrupt, ...] = ()
class PregelExecutableTask(NamedTuple):
libs/langgraph/tests/__snapshots__/test_pregel.ambr
+18
View File
@@ -682,6 +682,24 @@
'''
# ---
# name: test_dynamic_interrupt
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([__start__]):::first
tool_two_slow(tool_two_slow)
tool_two_fast(tool_two_fast)
__end__([__end__]):::last
__start__ -.-> tool_two_slow;
tool_two_slow --> __end__;
__start__ -.-> tool_two_fast;
tool_two_fast --> __end__;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
'''
# ---
# name: test_in_one_fan_out_state_graph_waiting_edge
'''
graph TD;
libs/langgraph/tests/test_pregel.py
+146 -30
View File
@@ -52,8 +52,8 @@ from langgraph.checkpoint.memory import MemorySaver
from langgraph.checkpoint.serde.base import SerializerProtocol
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.constants import ERROR, Send
from langgraph.errors import InvalidUpdateError
from langgraph.constants import ERROR, Interrupt, Send
from langgraph.errors import InvalidUpdateError, NodeInterrupt
from langgraph.graph import END, Graph
from langgraph.graph.graph import START
from langgraph.graph.message import MessageGraph, add_messages
@@ -2195,7 +2195,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
),
},
},
tasks=(PregelTask(AnyStr(), "tools"),),
tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -2239,7 +2239,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
"input": "what is weather in sf",
},
},
tasks=(PregelTask(AnyStr(), "tools"),),
tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -2400,7 +2400,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
),
},
},
tasks=(PregelTask(AnyStr(), "tools"),),
tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -2999,7 +2999,7 @@ def test_conditional_state_graph(
),
"intermediate_steps": [],
},
tasks=(PregelTask(AnyStr(), "tools"),),
tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3039,7 +3039,7 @@ def test_conditional_state_graph(
),
"intermediate_steps": [],
},
tasks=(PregelTask(AnyStr(), "tools"),),
tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3148,7 +3148,7 @@ def test_conditional_state_graph(
values={
"intermediate_steps": [],
},
tasks=(PregelTask(AnyStr(), "agent"),),
tasks=(PregelTask(AnyStr(), "agent", interrupts=(Interrupt("before"),)),),
next=("agent",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3173,7 +3173,7 @@ def test_conditional_state_graph(
),
"intermediate_steps": [],
},
tasks=(PregelTask(AnyStr(), "tools"),),
tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -3226,7 +3226,7 @@ def test_conditional_state_graph(
)
],
},
tasks=(PregelTask(AnyStr(), "agent"),),
tasks=(PregelTask(AnyStr(), "agent", interrupts=(Interrupt("before"),)),),
next=("agent",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -4828,7 +4828,13 @@ def test_message_graph(
id="ai1",
),
],
tasks=(PregelTask(AnyStr(), "tools"),),
tasks=(
PregelTask(
AnyStr(),
"tools",
interrupts=(Interrupt("before"),),
),
),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -4873,7 +4879,7 @@ def test_message_graph(
],
),
],
tasks=(PregelTask(AnyStr(), "tools"),),
tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -4955,7 +4961,7 @@ def test_message_graph(
id="ai2",
),
],
tasks=(PregelTask(AnyStr(), "tools"),),
tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -5547,7 +5553,13 @@ def test_root_graph(
id="ai1",
),
],
tasks=(PregelTask(AnyStr(), "tools"),),
tasks=(
PregelTask(
AnyStr(),
"tools",
interrupts=(Interrupt("before"),),
),
),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -5592,7 +5604,7 @@ def test_root_graph(
],
),
],
tasks=(PregelTask(AnyStr(), "tools"),),
tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -5674,7 +5686,7 @@ def test_root_graph(
id="ai2",
),
],
tasks=(PregelTask(AnyStr(), "tools"),),
tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
next=("tools",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -6120,6 +6132,72 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
]
def test_dynamic_interrupt(snapshot: SnapshotAssertion) -> None:
class State(TypedDict):
my_key: Annotated[str, operator.add]
market: str
def tool_two_node(s: State) -> State:
if s["market"] == "DE":
raise NodeInterrupt("Just because...")
return {"my_key": " all good"}
tool_two_graph = StateGraph(State)
tool_two_graph.add_node("tool_two", tool_two_node)
tool_two_graph.add_edge(START, "tool_two")
tool_two = tool_two_graph.compile()
assert tool_two.invoke({"my_key": "value", "market": "DE"}) == {
"my_key": "value",
"market": "DE",
}
assert tool_two.invoke({"my_key": "value", "market": "US"}) == {
"my_key": "value all good",
"market": "US",
}
with SqliteSaver.from_conn_string(":memory:") as saver:
tool_two = tool_two_graph.compile(checkpointer=saver)
# missing thread_id
with pytest.raises(ValueError, match="thread_id"):
tool_two.invoke({"my_key": "value", "market": "DE"})
thread1 = {"configurable": {"thread_id": "1"}}
# stop when about to enter node
assert tool_two.invoke({"my_key": "value ⛰️", "market": "DE"}, thread1) == {
"my_key": "value ⛰️",
"market": "DE",
}
assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [
{
"source": "loop",
"step": 0,
"writes": None,
},
{
"source": "input",
"step": -1,
"writes": {"my_key": "value ⛰️", "market": "DE"},
},
]
assert tool_two.get_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️", "market": "DE"},
next=("tool_two",),
tasks=(
PregelTask(
AnyStr(),
"tool_two",
interrupts=(Interrupt("during", "Just because..."),),
),
),
config=tool_two.checkpointer.get_tuple(thread1).config,
created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
metadata={"source": "loop", "step": 0, "writes": None},
parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config,
)
def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
class State(TypedDict):
my_key: Annotated[str, operator.add]
@@ -6172,7 +6250,13 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
]
assert tool_two.get_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️", "market": "DE"},
tasks=(PregelTask(AnyStr(), "tool_two_slow"),),
tasks=(
PregelTask(
AnyStr(),
"tool_two_slow",
interrupts=(Interrupt("before"),),
),
),
next=("tool_two_slow",),
config=tool_two.checkpointer.get_tuple(thread1).config,
created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
@@ -6206,7 +6290,13 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
}
assert tool_two.get_state(thread2) == StateSnapshot(
values={"my_key": "value", "market": "US"},
tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
tasks=(
PregelTask(
AnyStr(),
"tool_two_fast",
interrupts=(Interrupt("before"),),
),
),
next=("tool_two_fast",),
config=tool_two.checkpointer.get_tuple(thread2).config,
created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"],
@@ -6240,7 +6330,13 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
}
assert tool_two.get_state(thread3) == StateSnapshot(
values={"my_key": "value", "market": "US"},
tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
tasks=(
PregelTask(
AnyStr(),
"tool_two_fast",
interrupts=(Interrupt("before"),),
),
),
next=("tool_two_fast",),
config=tool_two.checkpointer.get_tuple(thread3).config,
created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"],
@@ -6251,7 +6347,13 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None:
tool_two.update_state(thread3, {"my_key": "key"}) # appends to my_key
assert tool_two.get_state(thread3) == StateSnapshot(
values={"my_key": "valuekey", "market": "US"},
tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
tasks=(
PregelTask(
AnyStr(),
"tool_two_fast",
interrupts=(Interrupt("before"),),
),
),
next=("tool_two_fast",),
config=tool_two.checkpointer.get_tuple(thread3).config,
created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"],
@@ -6344,7 +6446,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
"writes": {"my_key": "value", "market": "DE"},
},
"next": ["__start__"],
"tasks": [{"id": AnyStr(), "name": "__start__"}],
"tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}],
},
},
{
@@ -6373,7 +6475,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
"writes": None,
},
"next": ["prepare"],
"tasks": [{"id": AnyStr(), "name": "prepare"}],
"tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}],
},
},
{
@@ -6423,7 +6525,9 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
"writes": {"prepare": {"my_key": " prepared"}},
},
"next": ["tool_two_slow"],
"tasks": [{"id": AnyStr(), "name": "tool_two_slow"}],
"tasks": [
{"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()}
],
},
},
{
@@ -6473,7 +6577,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
"writes": {"tool_two_slow": {"my_key": " slow"}},
},
"next": ["finish"],
"tasks": [{"id": AnyStr(), "name": "finish"}],
"tasks": [{"id": AnyStr(), "name": "finish", "interrupts": ()}],
},
},
{
@@ -6544,7 +6648,13 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
}
assert tool_two.get_state(thread1) == StateSnapshot(
values={"my_key": "value prepared", "market": "DE"},
tasks=(PregelTask(AnyStr(), "tool_two_slow"),),
tasks=(
PregelTask(
AnyStr(),
"tool_two_slow",
interrupts=(Interrupt("before"),),
),
),
next=("tool_two_slow",),
config=tool_two.checkpointer.get_tuple(thread1).config,
created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
@@ -6582,7 +6692,13 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
}
assert tool_two.get_state(thread2) == StateSnapshot(
values={"my_key": "value prepared", "market": "US"},
tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
tasks=(
PregelTask(
AnyStr(),
"tool_two_fast",
interrupts=(Interrupt("before"),),
),
),
next=("tool_two_fast",),
config=tool_two.checkpointer.get_tuple(thread2).config,
created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"],
@@ -6629,7 +6745,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
"my_key": "value prepared slow",
"market": "DE",
},
tasks=(PregelTask(AnyStr(), "finish"),),
tasks=(PregelTask(AnyStr(), "finish", interrupts=(Interrupt("before"),)),),
next=("finish",),
config=tool_two.checkpointer.get_tuple(thread1).config,
created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
@@ -6648,7 +6764,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None:
"my_key": "value prepared slower",
"market": "DE",
},
tasks=(PregelTask(AnyStr(), "finish"),),
tasks=(PregelTask(AnyStr(), "finish", interrupts=(Interrupt("before"),)),),
next=("finish",),
config=tool_two.checkpointer.get_tuple(thread1).config,
created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
@@ -6905,7 +7021,7 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) ->
"query": "analyzed: query: what is weather in sf",
"docs": ["doc1", "doc2", "doc3", "doc4", "doc5"],
},
tasks=(PregelTask(AnyStr(), "qa"),),
tasks=(PregelTask(AnyStr(), "qa", interrupts=(Interrupt("before"),)),),
next=("qa",),
config=app_w_interrupt.checkpointer.get_tuple(config).config,
created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"],
@@ -7856,7 +7972,7 @@ def test_nested_graph_interrupts(
# test invoke w/ nested interrupt
config = {"configurable": {"thread_id": "1"}}
assert app.invoke({"my_key": "my value"}, config, debug=True) == {
assert app.invoke({"my_key": "my value"}, config) == {
"my_key": "hi my value",
}
assert list(app.get_state_history(config)) == [
libs/langgraph/tests/test_pregel_async.py
+132 -35
View File
@@ -47,8 +47,8 @@ from langgraph.checkpoint.base import (
)
from langgraph.checkpoint.memory import MemorySaver
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
from langgraph.constants import ERROR, Send
from langgraph.errors import InvalidUpdateError
from langgraph.constants import ERROR, Interrupt, Send
from langgraph.errors import InvalidUpdateError, NodeInterrupt
from langgraph.graph import END, Graph, StateGraph
from langgraph.graph.graph import START
from langgraph.graph.message import MessageGraph, add_messages
@@ -204,6 +204,61 @@ async def test_node_cancellation_on_other_node_exception() -> None:
assert inner_task_cancelled
@pytest.mark.parametrize(
"checkpointer_name",
["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"],
)
async def test_node_not_cancelled_on_other_node_interrupted(
checkpointer_name: str, request: pytest.FixtureRequest
) -> None:
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
class State(TypedDict):
hello: str
awhiles = 0
inner_task_cancelled = False
async def awhile(input: State) -> None:
nonlocal awhiles
awhiles += 1
try:
await asyncio.sleep(1)
return {"hello": "again"}
except asyncio.CancelledError:
nonlocal inner_task_cancelled
inner_task_cancelled = True
raise
async def iambad(input: State) -> None:
if input["hello"] != "bye":
raise NodeInterrupt("I am bad")
builder = StateGraph(State)
builder.add_node("agent", awhile)
builder.add_node("bad", iambad)
builder.set_conditional_entry_point(lambda _: ["agent", "bad"], then=END)
graph = builder.compile(checkpointer=checkpointer)
thread = {"configurable": {"thread_id": "1"}}
assert await graph.ainvoke({"hello": "world"}, thread) == {"hello": "world"}
assert not inner_task_cancelled
assert awhiles == 1
assert await graph.ainvoke(None, thread, debug=True) is None
assert not inner_task_cancelled
assert awhiles == 1
assert await graph.ainvoke({"hello": "bye"}, thread) == {"hello": "again"}
assert not inner_task_cancelled
assert awhiles == 2
async def test_step_timeout_on_stream_hang() -> None:
inner_task_cancelled = False
@@ -2366,7 +2421,13 @@ async def test_conditional_graph() -> None:
),
},
},
tasks=(PregelTask(AnyStr(), "tools"),),
tasks=(
PregelTask(
AnyStr(),
"tools",
interrupts=(Interrupt("before"),),
),
),
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[
@@ -2414,7 +2475,13 @@ async def test_conditional_graph() -> None:
"input": "what is weather in sf",
},
},
tasks=(PregelTask(AnyStr(), "tools"),),
tasks=(
PregelTask(
AnyStr(),
"tools",
interrupts=(Interrupt("before"),),
),
),
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[
@@ -2586,7 +2653,13 @@ async def test_conditional_graph() -> None:
),
},
},
tasks=(PregelTask(AnyStr(), "tools"),),
tasks=(
PregelTask(
AnyStr(),
"tools",
interrupts=(Interrupt("before"),),
),
),
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[
@@ -3152,7 +3225,13 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None:
),
"intermediate_steps": [],
},
tasks=(PregelTask(AnyStr(), "tools"),),
tasks=(
PregelTask(
AnyStr(),
"tools",
interrupts=(Interrupt("before"),),
),
),
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[
@@ -3196,7 +3275,7 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None:
),
"intermediate_steps": [],
},
tasks=(PregelTask(AnyStr(), "tools"),),
tasks=(PregelTask(AnyStr(), "tools", interrupts=(Interrupt("before"),)),),
next=("tools",),
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[
@@ -4746,7 +4825,13 @@ async def test_start_branch_then() -> None:
]
assert await tool_two.aget_state(thread1) == StateSnapshot(
values={"my_key": "value", "market": "DE"},
tasks=(PregelTask(AnyStr(), "tool_two_slow"),),
tasks=(
PregelTask(
AnyStr(),
"tool_two_slow",
interrupts=(Interrupt("before"),),
),
),
next=("tool_two_slow",),
config=(await tool_two.checkpointer.aget_tuple(thread1)).config,
created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[
@@ -4788,7 +4873,13 @@ async def test_start_branch_then() -> None:
}
assert await tool_two.aget_state(thread2) == StateSnapshot(
values={"my_key": "value", "market": "US"},
tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
tasks=(
PregelTask(
AnyStr(),
"tool_two_fast",
interrupts=(Interrupt("before"),),
),
),
next=("tool_two_fast",),
config=(await tool_two.checkpointer.aget_tuple(thread2)).config,
created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[
@@ -4830,7 +4921,13 @@ async def test_start_branch_then() -> None:
}
assert await tool_two.aget_state(thread3) == StateSnapshot(
values={"my_key": "value", "market": "US"},
tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
tasks=(
PregelTask(
AnyStr(),
"tool_two_fast",
interrupts=(Interrupt("before"),),
),
),
next=("tool_two_fast",),
config=(await tool_two.checkpointer.aget_tuple(thread3)).config,
created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[
@@ -4845,7 +4942,13 @@ async def test_start_branch_then() -> None:
await tool_two.aupdate_state(thread3, {"my_key": "key"}) # appends to my_key
assert await tool_two.aget_state(thread3) == StateSnapshot(
values={"my_key": "valuekey", "market": "US"},
tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
tasks=(
PregelTask(
AnyStr(),
"tool_two_fast",
interrupts=(Interrupt("before"),),
),
),
next=("tool_two_fast",),
config=(await tool_two.checkpointer.aget_tuple(thread3)).config,
created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[
@@ -4945,12 +5048,7 @@ async def test_branch_then() -> None:
"writes": {"my_key": "value", "market": "DE"},
},
"next": ["__start__"],
"tasks": [
{
"id": AnyStr(),
"name": "__start__",
}
],
"tasks": [{"id": AnyStr(), "name": "__start__", "interrupts": ()}],
},
},
{
@@ -4979,12 +5077,7 @@ async def test_branch_then() -> None:
"writes": None,
},
"next": ["prepare"],
"tasks": [
{
"id": AnyStr(),
"name": "prepare",
}
],
"tasks": [{"id": AnyStr(), "name": "prepare", "interrupts": ()}],
},
},
{
@@ -5035,10 +5128,7 @@ async def test_branch_then() -> None:
},
"next": ["tool_two_slow"],
"tasks": [
{
"id": AnyStr(),
"name": "tool_two_slow",
}
{"id": AnyStr(), "name": "tool_two_slow", "interrupts": ()}
],
},
},
@@ -5089,12 +5179,7 @@ async def test_branch_then() -> None:
"writes": {"tool_two_slow": {"my_key": " slow"}},
},
"next": ["finish"],
"tasks": [
{
"id": AnyStr(),
"name": "finish",
}
],
"tasks": [{"id": AnyStr(), "name": "finish", "interrupts": ()}],
},
},
{
@@ -5165,7 +5250,13 @@ async def test_branch_then() -> None:
}
assert await tool_two.aget_state(thread1) == StateSnapshot(
values={"my_key": "value prepared", "market": "DE"},
tasks=(PregelTask(AnyStr(), "tool_two_slow"),),
tasks=(
PregelTask(
AnyStr(),
"tool_two_slow",
interrupts=(Interrupt("before"),),
),
),
next=("tool_two_slow",),
config=(await tool_two.checkpointer.aget_tuple(thread1)).config,
created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[
@@ -5211,7 +5302,13 @@ async def test_branch_then() -> None:
}
assert await tool_two.aget_state(thread2) == StateSnapshot(
values={"my_key": "value prepared", "market": "US"},
tasks=(PregelTask(AnyStr(), "tool_two_fast"),),
tasks=(
PregelTask(
AnyStr(),
"tool_two_fast",
interrupts=(Interrupt("before"),),
),
),
next=("tool_two_fast",),
config=(await tool_two.checkpointer.aget_tuple(thread2)).config,
created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[