Update core and add tests for cancellation of astream and astream_events v2

2026-08-28 18:59:42 +02:00 · 2024-06-13 16:12:01 -07:00
parent c227cc0d94
commit e88e246a73
4 changed files with 181 additions and 9 deletions
@@ -155,7 +155,14 @@ class StateGraph(Graph):
    ) -> None:
        if not isinstance(node, str):
            action = node
-            node = getattr(action, "name", action.__name__)
+            if isinstance(action, Runnable):
+                node = action.name
+            else:
+                node = getattr(action, "__name__", action.__class__.__name__)
+            if node is None:
+                raise ValueError(
+                    "Node name must be provided if action is not a function"
+                )
        if node in self.channels:
            raise ValueError(f"'{node}' is already being used as a state key")
        return super().add_node(node, action)
@@ -1360,7 +1360,7 @@ class Pregel(
            except NameError:
                pass
            # wait for all background tasks to finish
-            await asyncio.gather(*bg)
+            await asyncio.shield(asyncio.gather(*bg))

    def invoke(
        self,
@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand.
+# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand.

 [[package]]
 name = "aiohttp"
@@ -1742,19 +1742,19 @@ langchain-core = ">=0.2.2rc1,<0.3"

 [[package]]
 name = "langchain-core"
-version = "0.2.5"
+version = "0.2.6"
 description = "Building applications with LLMs through composability"
 optional = false
 python-versions = "<4.0,>=3.8.1"
 files = [
-    {file = "langchain_core-0.2.5-py3-none-any.whl", hash = "sha256:abe5138f22acff23a079ec538be5268bbf97cf023d51987a0dd474d2a16cae3e"},
-    {file = "langchain_core-0.2.5.tar.gz", hash = "sha256:4a5c2f56b22396a63ef4790043660e393adbfa6832b978f023ca996a04b8e752"},
+    {file = "langchain_core-0.2.6-py3-none-any.whl", hash = "sha256:90521c9fc95d8f925e0d2e2d952382676aea6d3f8de611eda1b1810874c31e5d"},
+    {file = "langchain_core-0.2.6.tar.gz", hash = "sha256:9f0e38da722a558a6e95b6d86de01bd92e84558c47ac8ba599f02eab70a1c873"},
 ]

 [package.dependencies]
 jsonpatch = ">=1.33,<2.0"
-langsmith = ">=0.1.66,<0.2.0"
-packaging = ">=23.2,<24.0"
+langsmith = ">=0.1.75,<0.2.0"
+packaging = ">=23.2,<25"
 pydantic = ">=1,<3"
 PyYAML = ">=5.3"
 tenacity = ">=8.1.0,<9.0.0"
@@ -2865,6 +2865,7 @@ files = [
    {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:69b023b2b4daa7548bcfbd4aa3da05b3a74b772db9e23b982788168117739938"},
    {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:81e0b275a9ecc9c0c0c07b4b90ba548307583c125f54d5b6946cfee6360c733d"},
    {file = "PyYAML-6.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba336e390cd8e4d1739f42dfe9bb83a3cc2e80f567d8805e11b46f4a943f5515"},
+    {file = "PyYAML-6.0.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:326c013efe8048858a6d312ddd31d56e468118ad4cdeda36c719bf5bb6192290"},
    {file = "PyYAML-6.0.1-cp310-cp310-win32.whl", hash = "sha256:bd4af7373a854424dabd882decdc5579653d7868b8fb26dc7d0e99f823aa5924"},
    {file = "PyYAML-6.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:fd1592b3fdf65fff2ad0004b5e363300ef59ced41c2e6b3a99d4089fa8c5435d"},
    {file = "PyYAML-6.0.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:6965a7bc3cf88e5a1c3bd2e0b5c22f8d677dc88a455344035f03399034eb3007"},
@@ -2872,8 +2873,16 @@ files = [
    {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:42f8152b8dbc4fe7d96729ec2b99c7097d656dc1213a3229ca5383f973a5ed6d"},
    {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:062582fca9fabdd2c8b54a3ef1c978d786e0f6b3a1510e0ac93ef59e0ddae2bc"},
    {file = "PyYAML-6.0.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d2b04aac4d386b172d5b9692e2d2da8de7bfb6c387fa4f801fbf6fb2e6ba4673"},
+    {file = "PyYAML-6.0.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:e7d73685e87afe9f3b36c799222440d6cf362062f78be1013661b00c5c6f678b"},
    {file = "PyYAML-6.0.1-cp311-cp311-win32.whl", hash = "sha256:1635fd110e8d85d55237ab316b5b011de701ea0f29d07611174a1b42f1444741"},
    {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"},
+    {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"},
+    {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"},
+    {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"},
+    {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"},
+    {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"},
+    {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"},
+    {file = "PyYAML-6.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:0d3304d8c0adc42be59c5f8a4d9e3d7379e6955ad754aa9d6ab7a398b59dd1df"},
    {file = "PyYAML-6.0.1-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:50550eb667afee136e9a77d6dc71ae76a44df8b3e51e41b77f6de2932bfe0f47"},
    {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1fe35611261b29bd1de0070f0b2f47cb6ff71fa6595c077e42bd0c419fa27b98"},
    {file = "PyYAML-6.0.1-cp36-cp36m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:704219a11b772aea0d8ecd7058d0082713c3562b4e271b849ad7dc4a5c90c13c"},
@@ -2890,6 +2899,7 @@ files = [
    {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a0cd17c15d3bb3fa06978b4e8958dcdc6e0174ccea823003a106c7d4d7899ac5"},
    {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:28c119d996beec18c05208a8bd78cbe4007878c6dd15091efb73a30e90539696"},
    {file = "PyYAML-6.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7e07cbde391ba96ab58e532ff4803f79c4129397514e1413a7dc761ccd755735"},
+    {file = "PyYAML-6.0.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:49a183be227561de579b4a36efbb21b3eab9651dd81b1858589f796549873dd6"},
    {file = "PyYAML-6.0.1-cp38-cp38-win32.whl", hash = "sha256:184c5108a2aca3c5b3d3bf9395d50893a7ab82a38004c8f61c258d4428e80206"},
    {file = "PyYAML-6.0.1-cp38-cp38-win_amd64.whl", hash = "sha256:1e2722cc9fbb45d9b87631ac70924c11d3a401b2d7f410cc0e3bbf249f2dca62"},
    {file = "PyYAML-6.0.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:9eb6caa9a297fc2c2fb8862bc5370d0303ddba53ba97e71f08023b6cd73d16a8"},
@@ -2897,6 +2907,7 @@ files = [
    {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5773183b6446b2c99bb77e77595dd486303b4faab2b086e7b17bc6bef28865f6"},
    {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b786eecbdf8499b9ca1d697215862083bd6d2a99965554781d0d8d1ad31e13a0"},
    {file = "PyYAML-6.0.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bc1bf2925a1ecd43da378f4db9e4f799775d6367bdb94671027b73b393a7c42c"},
+    {file = "PyYAML-6.0.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:04ac92ad1925b2cff1db0cfebffb6ffc43457495c9b3c39d3fcae417d7125dc5"},
    {file = "PyYAML-6.0.1-cp39-cp39-win32.whl", hash = "sha256:faca3bdcf85b2fc05d06ff3fbc1f83e1391b3e724afa3feba7d13eeab355484c"},
    {file = "PyYAML-6.0.1-cp39-cp39-win_amd64.whl", hash = "sha256:510c9deebc5c0225e8c96813043e62b680ba2f9c50a08d3724c7f28a747d1486"},
    {file = "PyYAML-6.0.1.tar.gz", hash = "sha256:bfdf460b1736c775f2ba9f6a92bca30bc2095067b8a9d77876d1fad6cc3b4a43"},
@@ -1,8 +1,9 @@
 import asyncio
 import json
 import operator
+import time
 from collections import Counter
-from contextlib import asynccontextmanager, contextmanager
+from contextlib import aclosing, asynccontextmanager, contextmanager
 from typing import (
    Annotated,
    Any,
@@ -31,6 +32,7 @@ from langgraph.channels.binop import BinaryOperatorAggregate
 from langgraph.channels.context import Context
 from langgraph.channels.last_value import LastValue
 from langgraph.channels.topic import Topic
+from langgraph.checkpoint import BaseCheckpointSaver
 from langgraph.checkpoint.aiosqlite import AsyncSqliteSaver
 from langgraph.constants import Send
 from langgraph.errors import InvalidUpdateError
@@ -135,6 +137,158 @@ async def test_step_timeout_on_stream_hang() -> None:
    assert inner_task_cancelled


+@pytest.mark.parametrize(
+    "checkpointer",
+    [
+        MemorySaverAssertImmutable(),
+        AsyncSqliteSaver.from_conn_string(":memory:"),
+        None,
+    ],
+)
+async def test_cancel_graph_astream(
+    checkpointer: Optional[BaseCheckpointSaver],
+) -> None:
+    try:
+
+        class State(TypedDict):
+            value: int
+
+        class AwhileMaker:
+            def __init__(self) -> None:
+                self.reset()
+
+            async def __call__(self, input: State) -> Any:
+                self.started = True
+                try:
+                    await asyncio.sleep(1.5)
+                except asyncio.CancelledError:
+                    self.cancelled = True
+                    raise
+
+            def reset(self):
+                self.started = False
+                self.cancelled = False
+
+        async def alittlewhile(input: State) -> None:
+            await asyncio.sleep(0.6)
+            return {"value": 2}
+
+        awhile = AwhileMaker()
+        builder = StateGraph(State)
+        builder.add_node("awhile", awhile)
+        builder.add_node(alittlewhile)
+        builder.add_edge(START, "alittlewhile")
+        builder.add_edge("alittlewhile", "awhile")
+        graph = builder.compile(checkpointer=checkpointer)
+
+        # test interrupting astream
+        thread1: RunnableConfig = {"configurable": {"thread_id": 1}}
+        async with aclosing(graph.astream({"value": 1}, thread1)) as stream:
+            async for chunk in stream:
+                assert chunk == {"alittlewhile": {"value": 2}}
+                break
+
+        # node "awhile" should never start
+        assert awhile.started is False
+
+        # checkpoint with output of "alittlewhile" should not be saved
+        if checkpointer is not None:
+            state = await graph.aget_state(thread1)
+            assert state is not None
+            assert state.values == {"value": 1}
+            assert state.next == ("alittlewhile",)
+            assert state.metadata == {"source": "loop", "step": 0, "writes": None}
+    finally:
+        if getattr(checkpointer, "__aexit__", None):
+            await checkpointer.__aexit__(None, None, None)
+
+
+@pytest.mark.parametrize(
+    "checkpointer",
+    [
+        MemorySaverAssertImmutable(),
+        AsyncSqliteSaver.from_conn_string(":memory:"),
+        None,
+    ],
+)
+async def test_cancel_graph_astream_events_v2(
+    checkpointer: Optional[BaseCheckpointSaver],
+) -> None:
+    try:
+
+        class State(TypedDict):
+            value: int
+
+        class AwhileMaker:
+            def __init__(self) -> None:
+                self.reset()
+
+            async def __call__(self, input: State) -> Any:
+                self.started = True
+                try:
+                    await asyncio.sleep(1.5)
+                except asyncio.CancelledError:
+                    self.cancelled = True
+                    raise
+
+            def reset(self):
+                self.started = False
+                self.cancelled = False
+
+        async def alittlewhile(input: State) -> None:
+            await asyncio.sleep(0.6)
+            return {"value": 2}
+
+        awhile = AwhileMaker()
+        anotherwhile = AwhileMaker()
+        builder = StateGraph(State)
+        builder.add_node(alittlewhile)
+        builder.add_node("awhile", awhile)
+        builder.add_node("anotherwhile", anotherwhile)
+        builder.add_edge(START, "alittlewhile")
+        builder.add_edge("alittlewhile", "awhile")
+        builder.add_edge("awhile", "anotherwhile")
+        graph = builder.compile(checkpointer=checkpointer)
+
+        # test interrupting astream_events v2
+        got_event = False
+        thread2: RunnableConfig = {"configurable": {"thread_id": 2}}
+        async with aclosing(
+            graph.astream_events({"value": 1}, thread2, version="v2")
+        ) as stream:
+            async for chunk in stream:
+                if chunk["event"] == "on_chain_stream" and not chunk["parent_ids"]:
+                    print(time.perf_counter(), "got event out here", chunk)
+                    got_event = True
+                    assert chunk["data"]["chunk"] == {"alittlewhile": {"value": 2}}
+                    break
+
+        # did break
+        assert got_event
+
+        # node "awhile" starts but is cancelled
+        assert awhile.started is True
+        assert awhile.cancelled is True
+
+        # node "anotherwhile" should never start
+        assert anotherwhile.started is False
+
+        # checkpoint with output of "alittlewhile" should not be saved
+        if checkpointer is not None:
+            state = await graph.aget_state(thread2)
+            assert state is not None
+            assert state.values == {"value": 2}
+            assert state.next == ("awhile",)
+            assert state.metadata == {
+                "source": "loop",
+                "step": 1,
+                "writes": {"alittlewhile": {"value": 2}},
+            }
+    finally:
+        if getattr(checkpointer, "__aexit__", None):
+            await checkpointer.__aexit__(None, None, None)
+
+
 async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
    add_one = mocker.Mock(side_effect=lambda x: x + 1)
    chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")