diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index d3fbf1023..363ff375b 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -10,6 +10,7 @@ from langgraph.constants import ( EMPTY_SEQ, ERROR, INTERRUPT, + MISSING, NULL_TASK_ID, RESUME, RETURN, @@ -173,7 +174,8 @@ def map_output_updates( return updated: list[tuple[str, Any]] = [] for task, writes in output_tasks: - if rtn := next((value for chan, value in writes if chan == RETURN), None): + rtn = next((value for chan, value in writes if chan == RETURN), MISSING) + if rtn is not MISSING: updated.append((task.name, rtn)) elif isinstance(output_channels, str): updated.extend( diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 4a7114e46..790d22a06 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -1,6 +1,5 @@ import asyncio import concurrent.futures -import threading import time from functools import partial from typing import ( @@ -22,9 +21,11 @@ from langchain_core.callbacks import Callbacks from langgraph.constants import ( CONF, CONFIG_KEY_CALL, + CONFIG_KEY_SCRATCHPAD, CONFIG_KEY_SEND, ERROR, INTERRUPT, + MISSING, NO_WRITES, PUSH, RESUME, @@ -70,8 +71,6 @@ class PregelRunner: retry_policy: Optional[RetryPolicy] = None, get_waiter: Optional[Callable[[], concurrent.futures.Future[None]]] = None, ) -> Iterator[None]: - locks: dict[str, threading.Lock] = {} - def writer( task: PregelExecutableTask, writes: Sequence[tuple[str, Any]], @@ -81,24 +80,19 @@ class PregelRunner: if all(w[0] != PUSH for w in writes): return task.config[CONF][CONFIG_KEY_SEND](writes) - if task.id not in locks: - locks[task.id] = threading.Lock() - with locks[task.id]: - prev_length = len(task.writes) - # delegate to the underlying writer - task.config[CONF][CONFIG_KEY_SEND](writes) - # confirm no other concurrent writes were added - assert len(task.writes) == prev_length + len(writes) # schedule PUSH tasks, collect futures + scratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD] + scratchpad.setdefault("call_counter", 0) rtn: dict[int, Optional[concurrent.futures.Future]] = {} - for idx, w in enumerate(writes, start=prev_length): + for idx, w in enumerate(writes): # bail if not a PUSH write if w[0] != PUSH: continue # schedule the next task, if the callback returns one - if next_task := self.schedule_task( - task, idx, calls[idx - prev_length] if calls else None - ): + wcall = calls[idx] if calls else None + cnt = scratchpad["call_counter"] + scratchpad["call_counter"] += 1 + if next_task := self.schedule_task(task, cnt, wcall): if fut := next( ( f @@ -109,13 +103,19 @@ class PregelRunner: ): # if the parent task was retried, # the next task might already be running - rtn[idx - prev_length] = fut + rtn[idx] = fut elif next_task.writes: # if it already ran, return the result fut = concurrent.futures.Future() - if val := next(v for c, v in next_task.writes if c == RETURN): + if ( + val := next( + (v for c, v in next_task.writes if c == RETURN), MISSING + ) + ) and val is not MISSING: fut.set_result(val) - elif exc := next(v for c, v in next_task.writes if c == ERROR): + elif exc := next( + (v for c, v in next_task.writes if c == ERROR), None + ): fut.set_exception( exc if isinstance(exc, BaseException) @@ -123,7 +123,7 @@ class PregelRunner: ) else: fut.set_result(None) - rtn[idx - prev_length] = fut + rtn[idx] = fut else: # schedule the next task fut = self.submit( @@ -141,7 +141,7 @@ class PregelRunner: ) fut.add_done_callback(partial(self.commit, next_task)) futures[fut] = next_task - rtn[idx - prev_length] = fut + rtn[idx] = fut return [rtn.get(i) for i in range(len(writes))] def call( @@ -189,6 +189,8 @@ class PregelRunner: raise if not futures: # maybe `t` schuduled another task return + else: + tasks = () # don't reschedule this task # add waiter task if requested if get_waiter is not None: futures[get_waiter()] = None @@ -255,8 +257,6 @@ class PregelRunner: retry_policy: Optional[RetryPolicy] = None, get_waiter: Optional[Callable[[], asyncio.Future[None]]] = None, ) -> AsyncIterator[None]: - locks: dict[str, threading.Lock] = {} - def writer( task: PregelExecutableTask, writes: Sequence[tuple[str, Any]], @@ -266,23 +266,19 @@ class PregelRunner: if all(w[0] != PUSH for w in writes): return task.config[CONF][CONFIG_KEY_SEND](writes) - if task.id not in locks: - locks[task.id] = threading.Lock() - with locks[task.id]: - prev_length = len(task.writes) - # delegate to the underlying writer - task.config[CONF][CONFIG_KEY_SEND](writes) - # confirm no other concurrent writes were added - assert len(task.writes) == prev_length + len(writes) # schedule PUSH tasks, collect futures + scratchpad = task.config[CONF][CONFIG_KEY_SCRATCHPAD] + scratchpad.setdefault("call_counter", 0) rtn: dict[int, Optional[asyncio.Future]] = {} - for idx, w in enumerate(writes, start=prev_length): + for idx, w in enumerate(writes): # bail if not a PUSH write if w[0] != PUSH: continue # schedule the next task, if the callback returns one - wcall = calls[idx - prev_length] if calls is not None else None - if next_task := self.schedule_task(task, idx, wcall): + wcall = calls[idx] if calls is not None else None + cnt = scratchpad["call_counter"] + scratchpad["call_counter"] += 1 + if next_task := self.schedule_task(task, cnt, wcall): # if the parent task was retried, # the next task might already be running if fut := next( @@ -295,13 +291,19 @@ class PregelRunner: ): # if the parent task was retried, # the next task might already be running - rtn[idx - prev_length] = fut + rtn[idx] = fut elif next_task.writes: # if it already ran, return the result fut = asyncio.Future() - if val := next(v for c, v in next_task.writes if c == RETURN): + if ( + val := next( + (v for c, v in next_task.writes if c == RETURN), MISSING + ) + ) and val is not MISSING: fut.set_result(val) - elif exc := next(v for c, v in next_task.writes if c == ERROR): + elif exc := next( + (v for c, v in next_task.writes if c == ERROR), None + ): fut.set_exception( exc if isinstance(exc, BaseException) @@ -309,7 +311,7 @@ class PregelRunner: ) else: fut.set_result(None) - rtn[idx - prev_length] = fut + rtn[idx] = fut else: # schedule the next task fut = cast( @@ -333,7 +335,7 @@ class PregelRunner: ) fut.add_done_callback(partial(self.commit, next_task)) futures[fut] = next_task - rtn[idx - prev_length] = fut + rtn[idx] = fut return [rtn.get(i) for i in range(len(writes))] def call( @@ -388,6 +390,8 @@ class PregelRunner: raise if not futures: # maybe `t` schuduled another task return + else: + tasks = () # don't reschedule this task # add waiter task if requested if get_waiter is not None: futures[get_waiter()] = None diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 066d7d106..0b9fb9b1b 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -342,9 +342,12 @@ class LoopProtocol: class PregelScratchpad(TypedDict, total=False): + # interrupt interrupt_counter: int used_null_resume: bool resume: list[Any] + # call + call_counter: int def interrupt(value: Any) -> Any: diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 710f222bb..12b7c6863 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -5260,3 +5260,60 @@ def test_multiple_updates() -> None: {"node_a": [{"foo": "a1"}, {"foo": "a2"}]}, {"node_b": {"foo": "b"}}, ] + + +def test_falsy_return_from_task() -> None: + """Test with a falsy return from a task.""" + checkpointer = MemorySaver() + + @task + def falsy_task() -> bool: + return False + + @entrypoint(checkpointer=checkpointer) + def graph(state: dict) -> dict: + """React tool.""" + falsy_task().result() + interrupt("test") + + configurable = {"configurable": {"thread_id": uuid.uuid4()}} + graph.invoke({"a": 5}, configurable) + graph.invoke(Command(resume="123"), configurable) + + +def test_multiple_interrupts_imperative() -> 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 + def double(x: int) -> int: + """Increment the counter.""" + nonlocal counter + counter += 1 + return 2 * x + + @entrypoint(checkpointer=checkpointer) + def graph(state: dict) -> dict: + """React tool.""" + + values = [] + + for idx in [1, 2, 3]: + values.extend([double(idx).result(), interrupt({"a": "boo"})]) + + return {"values": values} + + configurable = {"configurable": {"thread_id": uuid.uuid4()}} + graph.invoke({}, configurable) + graph.invoke(Command(resume="a"), configurable) + graph.invoke(Command(resume="b"), configurable) + result = graph.invoke(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 diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index cafd2fbb3..427ee4041 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6693,3 +6693,69 @@ async def test_multiple_updates() -> None: {"node_a": [{"foo": "a1"}, {"foo": "a2"}]}, {"node_b": {"foo": "b"}}, ] + + +@pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) +async def test_falsy_return_from_task() -> None: + """Test with a falsy return from a task.""" + checkpointer = MemorySaver() + + @task + async def falsy_task() -> bool: + return False + + @entrypoint(checkpointer=checkpointer) + async def graph(state: dict) -> dict: + """React tool.""" + await falsy_task() + interrupt("test") + + configurable = {"configurable": {"thread_id": uuid.uuid4()}} + await graph.ainvoke({"a": 5}, configurable) + await graph.ainvoke(Command(resume="123"), configurable) + + +@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.func import entrypoint, task + + counter = 0 + + @task + async def double(x: int) -> int: + """Increment the counter.""" + nonlocal counter + counter += 1 + return 2 * x + + async with awith_checkpointer(checkpointer_name) as checkpointer: + + @entrypoint(checkpointer=checkpointer) + async def graph(state: dict) -> dict: + """React tool.""" + + values = [] + + for idx in [1, 2, 3]: + values.extend([await double(idx), interrupt({"a": "boo"})]) + + return {"values": values} + + configurable = {"configurable": {"thread_id": str(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