From fe87c44da7d04b047b4c3b5655629cac5ff0ced5 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 14 Jun 2024 13:32:21 -0700 Subject: [PATCH 1/2] Add tests to ensure checkpointer errors are not swallowed --- poetry.lock | 6 ++-- tests/test_pregel.py | 46 ++++++++++++++++++++++++++ tests/test_pregel_async.py | 66 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 115 insertions(+), 3 deletions(-) diff --git a/poetry.lock b/poetry.lock index 4eab3cf22..6bf3e1244 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1742,13 +1742,13 @@ langchain-core = ">=0.2.2rc1,<0.3" [[package]] name = "langchain-core" -version = "0.2.6" +version = "0.2.7" description = "Building applications with LLMs through composability" optional = false python-versions = "<4.0,>=3.8.1" files = [ - {file = "langchain_core-0.2.6-py3-none-any.whl", hash = "sha256:90521c9fc95d8f925e0d2e2d952382676aea6d3f8de611eda1b1810874c31e5d"}, - {file = "langchain_core-0.2.6.tar.gz", hash = "sha256:9f0e38da722a558a6e95b6d86de01bd92e84558c47ac8ba599f02eab70a1c873"}, + {file = "langchain_core-0.2.7-py3-none-any.whl", hash = "sha256:fd02e153c898486dd728d634684ffc64bc257ff2ba443dc7e53d017ac0bf4658"}, + {file = "langchain_core-0.2.7.tar.gz", hash = "sha256:b0b1b6dfbdedb39426fcb8bd3f07e40eec7964856e3fc384c420ca6dba61b34e"}, ] [package.dependencies] diff --git a/tests/test_pregel.py b/tests/test_pregel.py index b6f06eac1..80c47840b 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -30,10 +30,17 @@ from langsmith import traceable from pytest_mock import MockerFixture from syrupy import SnapshotAssertion +from langgraph.channels.base import BaseChannel 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.base import ( + Checkpoint, + CheckpointMetadata, + CheckpointTuple, +) +from langgraph.checkpoint.memory import MemorySaver from langgraph.checkpoint.sqlite import SqliteSaver from langgraph.constants import Send from langgraph.errors import InvalidUpdateError @@ -171,6 +178,45 @@ def test_graph_validation() -> None: graph.compile() +def test_checkpoint_errors() -> None: + class FaultyGetCheckpointer(MemorySaver): + def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None: + raise ValueError("Faulty get_tuple") + + class FaultyPutCheckpointer(MemorySaver): + def put( + self, + config: RunnableConfig, + checkpoint: Checkpoint, + metadata: CheckpointMetadata, + ) -> RunnableConfig: + raise ValueError("Faulty put") + + class FaultyVersionCheckpointer(MemorySaver): + def get_next_version(self, current: Optional[int], channel: BaseChannel) -> int: + raise ValueError("Faulty get_next_version") + + def logic(inp: str) -> str: + return "" + + builder = Graph() + builder.add_node("agent", logic) + builder.set_entry_point("agent") + builder.set_finish_point("agent") + + graph = builder.compile(checkpointer=FaultyGetCheckpointer()) + with pytest.raises(ValueError, match="Faulty get_tuple"): + graph.invoke("", {"configurable": {"thread_id": "thread-1"}}) + + graph = builder.compile(checkpointer=FaultyPutCheckpointer()) + with pytest.raises(ValueError, match="Faulty put"): + graph.invoke("", {"configurable": {"thread_id": "thread-1"}}) + + graph = builder.compile(checkpointer=FaultyVersionCheckpointer()) + with pytest.raises(ValueError, match="Faulty get_next_version"): + graph.invoke("", {"configurable": {"thread_id": "thread-1"}}) + + def test_reducer_before_first_node() -> None: from langchain_core.messages import HumanMessage diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 273f16d2c..e5a867374 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -29,12 +29,15 @@ from langchain_core.utils.aiter import aclosing from pytest_mock import MockerFixture from syrupy import SnapshotAssertion +from langgraph.channels.base import BaseChannel 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.checkpoint.base import Checkpoint, CheckpointMetadata, CheckpointTuple +from langgraph.checkpoint.memory import MemorySaver from langgraph.constants import Send from langgraph.errors import InvalidUpdateError from langgraph.graph import END, Graph, StateGraph @@ -56,6 +59,69 @@ from tests.memory_assert import ( ) +async def test_checkpoint_errors() -> None: + class FaultyGetCheckpointer(MemorySaver): + async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None: + raise ValueError("Faulty get_tuple") + + class FaultyPutCheckpointer(MemorySaver): + async def aput( + self, + config: RunnableConfig, + checkpoint: Checkpoint, + metadata: CheckpointMetadata, + ) -> RunnableConfig: + raise ValueError("Faulty put") + + class FaultyVersionCheckpointer(MemorySaver): + def get_next_version(self, current: Optional[int], channel: BaseChannel) -> int: + raise ValueError("Faulty get_next_version") + + def logic(inp: str) -> str: + return "" + + builder = Graph() + builder.add_node("agent", logic) + builder.set_entry_point("agent") + builder.set_finish_point("agent") + + graph = builder.compile(checkpointer=FaultyGetCheckpointer()) + with pytest.raises(ValueError, match="Faulty get_tuple"): + await graph.ainvoke("", {"configurable": {"thread_id": "thread-1"}}) + with pytest.raises(ValueError, match="Faulty get_tuple"): + async for _ in graph.astream("", {"configurable": {"thread_id": "thread-2"}}): + pass + with pytest.raises(ValueError, match="Faulty get_tuple"): + async for _ in graph.astream_events( + "", {"configurable": {"thread_id": "thread-3"}}, version="v2" + ): + pass + + graph = builder.compile(checkpointer=FaultyPutCheckpointer()) + with pytest.raises(ValueError, match="Faulty put"): + await graph.ainvoke("", {"configurable": {"thread_id": "thread-1"}}) + with pytest.raises(ValueError, match="Faulty put"): + async for _ in graph.astream("", {"configurable": {"thread_id": "thread-2"}}): + pass + with pytest.raises(ValueError, match="Faulty put"): + async for _ in graph.astream_events( + "", {"configurable": {"thread_id": "thread-3"}}, version="v2" + ): + pass + + graph = builder.compile(checkpointer=FaultyVersionCheckpointer()) + with pytest.raises(ValueError, match="Faulty get_next_version"): + await graph.ainvoke("", {"configurable": {"thread_id": "thread-1"}}) + with pytest.raises(ValueError, match="Faulty get_next_version"): + async for _ in graph.astream("", {"configurable": {"thread_id": "thread-2"}}): + pass + with pytest.raises(ValueError, match="Faulty get_next_version"): + async for _ in graph.astream_events( + "", {"configurable": {"thread_id": "thread-3"}}, version="v2" + ): + pass + + async def test_node_cancellation_on_external_cancel() -> None: inner_task_cancelled = False From e8dc9f0678bb29795757db87d57088a5bcef42e6 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 14 Jun 2024 13:35:28 -0700 Subject: [PATCH 2/2] old py --- tests/test_pregel.py | 2 +- tests/test_pregel_async.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 80c47840b..55953a18f 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -180,7 +180,7 @@ def test_graph_validation() -> None: def test_checkpoint_errors() -> None: class FaultyGetCheckpointer(MemorySaver): - def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None: + def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: raise ValueError("Faulty get_tuple") class FaultyPutCheckpointer(MemorySaver): diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index e5a867374..b5cfb3a80 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -61,7 +61,7 @@ from tests.memory_assert import ( async def test_checkpoint_errors() -> None: class FaultyGetCheckpointer(MemorySaver): - async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None: + async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: raise ValueError("Faulty get_tuple") class FaultyPutCheckpointer(MemorySaver):