This commit is contained in:
Nuno Campos
2025-01-15 10:21:40 -08:00
parent 1e9a372dd7
commit a7bb96da98
+25 -20
View File
@@ -6714,12 +6714,15 @@ async def test_falsy_return_from_task() -> None:
await graph.ainvoke(Command(resume="123"), configurable)
async def test_multiple_interrupts_imperative() -> None:
@pytest.mark.skipif(
sys.version_info < (3, 11),
reason="Python 3.11+ is required for async contextvars support",
)
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_multiple_interrupts_imperative(checkpointer_name: str) -> None:
"""Test multiple interrupts with an imperative API."""
from langgraph.checkpoint.memory import MemorySaver
from langgraph.func import entrypoint, task
checkpointer = MemorySaver()
counter = 0
@task
@@ -6729,24 +6732,26 @@ async def test_multiple_interrupts_imperative() -> None:
counter += 1
return 2 * x
@entrypoint(checkpointer=checkpointer)
async def graph(state: dict) -> dict:
"""React tool."""
async with awith_checkpointer(checkpointer_name) as checkpointer:
values = []
@entrypoint(checkpointer=checkpointer)
async def graph(state: dict) -> dict:
"""React tool."""
for idx in [1, 2, 3]:
values.extend([await double(idx), interrupt({"a": "boo"})])
values = []
return {"values": values}
for idx in [1, 2, 3]:
values.extend([await double(idx), interrupt({"a": "boo"})])
configurable = {"configurable": {"thread_id": uuid.uuid4()}}
await graph.ainvoke({}, configurable)
await graph.ainvoke(Command(resume="a"), configurable)
await graph.ainvoke(Command(resume="b"), configurable)
result = await graph.ainvoke(Command(resume="c"), configurable)
# `double` value should be cached appropriately when used w/ `interrupt`
assert result == {
"values": [2, "a", 4, "b", 6, "c"],
}
assert counter == 3
return {"values": values}
configurable = {"configurable": {"thread_id": uuid.uuid4()}}
await graph.ainvoke({}, configurable)
await graph.ainvoke(Command(resume="a"), configurable)
await graph.ainvoke(Command(resume="b"), configurable)
result = await graph.ainvoke(Command(resume="c"), configurable)
# `double` value should be cached appropriately when used w/ `interrupt`
assert result == {
"values": [2, "a", 4, "b", 6, "c"],
}
assert counter == 3