From bc17c3522b465467119a0d456575d819a98cd4a0 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sat, 24 May 2025 14:27:13 -0700 Subject: [PATCH] Remove old checkpoint test fixtures - Now all tests fully migrated to more recent sync_checkpointer and async_checkpointer fixtures for parametrising on checkpointer - Use sync/async_store fixtures where tests used only in memory store - Remove unused "should snapshot" check for older versions of langchain core no longer tested against --- libs/langgraph/.claude/settings.local.json | 5 +- .../tests/__snapshots__/test_large_cases.ambr | 2 +- .../tests/__snapshots__/test_pregel.ambr | 70 - libs/langgraph/tests/conftest.py | 96 - libs/langgraph/tests/conftest_checkpointer.py | 5 +- libs/langgraph/tests/test_large_cases.py | 494 +- .../langgraph/tests/test_large_cases_async.py | 9099 ++++++++--------- libs/langgraph/tests/test_pregel.py | 583 +- libs/langgraph/tests/test_pregel_async.py | 5841 ++++++----- libs/prebuilt/tests/conftest.py | 7 - libs/prebuilt/tests/test_react_agent.py | 17 - libs/prebuilt/tests/test_tool_node.py | 9 - 12 files changed, 7817 insertions(+), 8411 deletions(-) diff --git a/libs/langgraph/.claude/settings.local.json b/libs/langgraph/.claude/settings.local.json index 67e11f281..d4dec4ec5 100644 --- a/libs/langgraph/.claude/settings.local.json +++ b/libs/langgraph/.claude/settings.local.json @@ -2,7 +2,10 @@ "permissions": { "allow": [ "Bash(rg:*)", - "Bash(python:*)" + "Bash(python:*)", + "Bash(grep:*)", + "Bash(sed:*)", + "Bash(awk:*)" ], "deny": [] } diff --git a/libs/langgraph/tests/__snapshots__/test_large_cases.ambr b/libs/langgraph/tests/__snapshots__/test_large_cases.ambr index cdd896b8d..34c117a1b 100644 --- a/libs/langgraph/tests/__snapshots__/test_large_cases.ambr +++ b/libs/langgraph/tests/__snapshots__/test_large_cases.ambr @@ -422,7 +422,7 @@ ''' # --- -# name: test_start_branch_then[memory] +# name: test_start_branch_then[memory-in_memory] ''' --- config: diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index 9559651f8..713a9d351 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -569,76 +569,6 @@ ''' # --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[memory] - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query -.-> retriever_two; - qa --> __end__; - - ''' -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[memory].1 - dict({ - 'definitions': dict({ - 'InnerObject': dict({ - 'properties': dict({ - 'yo': dict({ - 'title': 'Yo', - 'type': 'integer', - }), - }), - 'required': list([ - 'yo', - ]), - 'title': 'InnerObject', - 'type': 'object', - }), - }), - 'properties': dict({ - 'inner': dict({ - '$ref': '#/definitions/InnerObject', - }), - 'query': dict({ - 'title': 'Query', - 'type': 'string', - }), - }), - 'required': list([ - 'query', - 'inner', - ]), - 'title': 'Input', - 'type': 'object', - }) -# --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[memory].2 - dict({ - 'properties': dict({ - 'answer': dict({ - 'title': 'Answer', - 'type': 'string', - }), - 'docs': dict({ - 'items': dict({ - 'type': 'string', - }), - 'title': 'Docs', - 'type': 'array', - }), - }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', - 'type': 'object', - }) -# --- # name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[memory] ''' graph TD; diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index 75f83cbdc..c6b640baa 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -1,11 +1,7 @@ from collections.abc import AsyncIterator, Iterator -from contextlib import asynccontextmanager -from typing import Optional from uuid import UUID import pytest -from langchain_core import __version__ as core_version -from packaging import version from pytest_mock import MockerFixture from langgraph.cache.base import BaseCache @@ -35,14 +31,6 @@ from tests.conftest_store import ( _store_postgres_pool, ) -pytest.register_assert_rewrite("tests.memory_assert") - -# TODO: fix this once core is released -IS_LANGCHAIN_CORE_030_OR_GREATER = version.parse(core_version) >= version.parse( - "0.3.0.dev0" -) -SHOULD_CHECK_SNAPSHOTS = IS_LANGCHAIN_CORE_030_OR_GREATER - @pytest.fixture def anyio_backend(): @@ -62,73 +50,6 @@ def checkpoint_during(request: pytest.FixtureRequest) -> bool: return request.param -# --- start of deprecated fixtures --- - - -@pytest.fixture(scope="function") -def checkpointer_memory(): - with _checkpointer_memory() as checkpointer: - yield checkpointer - - -@pytest.fixture(scope="function") -def checkpointer_sqlite(): - with _checkpointer_sqlite() as checkpointer: - yield checkpointer - - -@pytest.fixture(scope="function") -def checkpointer_sqlite_aes(): - with _checkpointer_sqlite_aes() as checkpointer: - yield checkpointer - - -@pytest.fixture(scope="function") -def checkpointer_postgres(): - with _checkpointer_postgres() as checkpointer: - yield checkpointer - - -@pytest.fixture(scope="function") -def checkpointer_postgres_pipe(): - with _checkpointer_postgres_pipe() as checkpointer: - yield checkpointer - - -@pytest.fixture(scope="function") -def checkpointer_postgres_pool(): - with _checkpointer_postgres_pool() as checkpointer: - yield checkpointer - - -@asynccontextmanager -async def awith_checkpointer( - checkpointer_name: Optional[str], -) -> AsyncIterator[BaseCheckpointSaver]: - if checkpointer_name is None: - yield None - elif checkpointer_name == "memory": - with _checkpointer_memory() as checkpointer: - yield checkpointer - elif checkpointer_name == "sqlite_aio": - async with _checkpointer_sqlite_aio() as checkpointer: - yield checkpointer - elif checkpointer_name == "postgres_aio": - async with _checkpointer_postgres_aio() as checkpointer: - yield checkpointer - elif checkpointer_name == "postgres_aio_pipe": - async with _checkpointer_postgres_aio_pipe() as checkpointer: - yield checkpointer - elif checkpointer_name == "postgres_aio_pool": - async with _checkpointer_postgres_aio_pool() as checkpointer: - yield checkpointer - else: - raise NotImplementedError(f"Unknown checkpointer: {checkpointer_name}") - - -# --- end of deprecated fixtures --- - - @pytest.fixture(scope="function", params=["sqlite", "memory"]) def cache(request: pytest.FixtureRequest) -> Iterator[BaseCache]: if request.param == "sqlite": @@ -255,20 +176,3 @@ async def async_checkpointer( yield checkpointer else: raise NotImplementedError(f"Unknown checkpointer: {checkpointer_name}") - - -ALL_CHECKPOINTERS_SYNC = [ - "memory", - "sqlite", - "sqlite_aes", - "postgres", - "postgres_pipe", - "postgres_pool", -] -ALL_CHECKPOINTERS_ASYNC = [ - "memory", - "sqlite_aio", - "postgres_aio", - "postgres_aio_pipe", - "postgres_aio_pool", -] diff --git a/libs/langgraph/tests/conftest_checkpointer.py b/libs/langgraph/tests/conftest_checkpointer.py index 84bd16a1c..ba15a8251 100644 --- a/libs/langgraph/tests/conftest_checkpointer.py +++ b/libs/langgraph/tests/conftest_checkpointer.py @@ -11,7 +11,10 @@ from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver from langgraph.checkpoint.serde.encrypted import EncryptedSerializer from langgraph.checkpoint.sqlite import SqliteSaver from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver -from tests.memory_assert import MemorySaverAssertImmutable + +pytest.register_assert_rewrite("tests.memory_assert") + +from tests.memory_assert import MemorySaverAssertImmutable # noqa: E402 DEFAULT_POSTGRES_URI = "postgres://postgres:postgres@localhost:5442/" diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index 5e56d0062..b570f086a 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -14,6 +14,7 @@ from typing_extensions import TypedDict from langgraph.channels.last_value import LastValue from langgraph.channels.untracked_value import UntrackedValue from langgraph.checkpoint.base import BaseCheckpointSaver +from langgraph.checkpoint.memory import InMemorySaver from langgraph.constants import END, PULL, PUSH, START from langgraph.errors import NodeInterrupt from langgraph.graph import StateGraph @@ -22,6 +23,7 @@ from langgraph.graph.message import MessageGraph, MessagesState, add_messages from langgraph.prebuilt.chat_agent_executor import create_react_agent from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import Channel, Pregel +from langgraph.store.base import BaseStore from langgraph.store.memory import InMemoryStore from langgraph.types import ( Command, @@ -36,7 +38,6 @@ from langgraph.types import ( from tests.agents import AgentAction, AgentFinish from tests.any_int import AnyInt from tests.any_str import AnyDict, AnyStr, UnsortedSequence -from tests.conftest import ALL_CHECKPOINTERS_SYNC, SHOULD_CHECK_SNAPSHOTS from tests.fake_chat import FakeChatModel from tests.fake_tracer import FakeTracer from tests.messages import ( @@ -47,11 +48,9 @@ from tests.messages import ( ) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_invoke_two_processes_in_out_interrupt( - request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture + sync_checkpointer: BaseCheckpointSaver, mocker: MockerFixture ) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") @@ -65,7 +64,7 @@ def test_invoke_two_processes_in_out_interrupt( }, input_channels="input", output_channels="output", - checkpointer=checkpointer, + checkpointer=sync_checkpointer, interrupt_after_nodes=["one"], ) thread1 = {"configurable": {"thread_id": "1"}} @@ -75,7 +74,7 @@ def test_invoke_two_processes_in_out_interrupt( assert app.invoke(2, thread1) is None # inbox == 3 - checkpoint = checkpointer.get(thread1) + checkpoint = sync_checkpointer.get(thread1) assert checkpoint is not None assert checkpoint["channel_values"]["inbox"] == 3 @@ -86,7 +85,7 @@ def test_invoke_two_processes_in_out_interrupt( assert app.invoke(20, thread1) is None # inbox == 21 - checkpoint = checkpointer.get(thread1) + checkpoint = sync_checkpointer.get(thread1) assert checkpoint is not None assert checkpoint["channel_values"]["inbox"] == 21 @@ -110,7 +109,6 @@ def test_invoke_two_processes_in_out_interrupt( snapshot = app.get_state(thread2) assert snapshot.next == () - # list history history = [c for c in app.get_state_history(thread1)] assert history == [ @@ -507,19 +505,14 @@ def test_fork_always_re_runs_nodes( ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_conditional_graph( - snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str + snapshot: SnapshotAssertion, sync_checkpointer: BaseCheckpointSaver ) -> None: from langchain_core.language_models.fake import FakeStreamingListLLM from langchain_core.prompts import PromptTemplate from langchain_core.runnables import RunnablePassthrough from langchain_core.tools import tool - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - # Assemble the tools @tool() def search_api(query: str) -> str: @@ -591,7 +584,7 @@ def test_conditional_graph( app = workflow.compile() - if SHOULD_CHECK_SNAPSHOTS and checkpointer_name == "memory": + if isinstance(sync_checkpointer, InMemorySaver): assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot assert app.get_graph().draw_mermaid(with_styles=False) == snapshot assert app.get_graph().draw_mermaid() == snapshot @@ -719,7 +712,7 @@ def test_conditional_graph( # test state get/update methods with interrupt_after app_w_interrupt = workflow.compile( - checkpointer=checkpointer, + checkpointer=sync_checkpointer, interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -774,7 +767,8 @@ def test_conditional_graph( }, "thread_id": "1", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -834,7 +828,8 @@ def test_conditional_graph( }, "thread_id": "1", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -963,7 +958,8 @@ def test_conditional_graph( }, "thread_id": "1", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -971,7 +967,7 @@ def test_conditional_graph( # test state get/update methods with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=checkpointer, + checkpointer=sync_checkpointer, interrupt_before=["tools"], ) config = {"configurable": {"thread_id": "2"}} @@ -1027,7 +1023,8 @@ def test_conditional_graph( }, "thread_id": "2", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -1081,7 +1078,8 @@ def test_conditional_graph( }, "thread_id": "2", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -1210,7 +1208,8 @@ def test_conditional_graph( }, "thread_id": "2", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -1218,7 +1217,7 @@ def test_conditional_graph( # test re-invoke to continue with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=checkpointer, + checkpointer=sync_checkpointer, interrupt_before=["tools"], ) config = {"configurable": {"thread_id": "3"}} @@ -1274,7 +1273,8 @@ def test_conditional_graph( }, "thread_id": "3", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -1398,21 +1398,14 @@ def test_conditional_graph( ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_conditional_state_graph( snapshot: SnapshotAssertion, - mocker: MockerFixture, - request: pytest.FixtureRequest, - checkpointer_name: str, + sync_checkpointer: BaseCheckpointSaver, ) -> None: from langchain_core.language_models.fake import FakeStreamingListLLM from langchain_core.prompts import PromptTemplate from langchain_core.tools import tool - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - class AgentState(TypedDict, total=False): input: Annotated[str, UntrackedValue] agent_outcome: Optional[Union[AgentAction, AgentFinish]] @@ -1495,7 +1488,7 @@ def test_conditional_state_graph( app = workflow.compile() - if SHOULD_CHECK_SNAPSHOTS and checkpointer_name == "memory": + if isinstance(sync_checkpointer, InMemorySaver): assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot @@ -1585,7 +1578,7 @@ def test_conditional_state_graph( # test state get/update methods with interrupt_after app_w_interrupt = workflow.compile( - checkpointer=checkpointer, + checkpointer=sync_checkpointer, interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -1637,7 +1630,8 @@ def test_conditional_state_graph( }, "thread_id": "1", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -1687,7 +1681,8 @@ def test_conditional_state_graph( }, "thread_id": "1", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -1770,7 +1765,8 @@ def test_conditional_state_graph( }, "thread_id": "1", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -1778,7 +1774,7 @@ def test_conditional_state_graph( # test state get/update methods with interrupt_before app_w_interrupt = workflow.compile( - checkpointer=checkpointer, + checkpointer=sync_checkpointer, interrupt_before=["tools"], debug=True, ) @@ -1830,7 +1826,8 @@ def test_conditional_state_graph( }, "thread_id": "2", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -1880,7 +1877,8 @@ def test_conditional_state_graph( }, "thread_id": "2", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -1963,14 +1961,15 @@ def test_conditional_state_graph( }, "thread_id": "2", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) # test w interrupt before all app_w_interrupt = workflow.compile( - checkpointer=checkpointer, + checkpointer=sync_checkpointer, interrupt_before="*", debug=True, ) @@ -2004,7 +2003,8 @@ def test_conditional_state_graph( "writes": None, "thread_id": "3", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -2052,7 +2052,8 @@ def test_conditional_state_graph( }, "thread_id": "3", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -2121,7 +2122,8 @@ def test_conditional_state_graph( }, "thread_id": "3", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -2141,7 +2143,7 @@ def test_conditional_state_graph( # test w interrupt after all app_w_interrupt = workflow.compile( - checkpointer=checkpointer, + checkpointer=sync_checkpointer, interrupt_after="*", ) config = {"configurable": {"thread_id": "4"}} @@ -2192,7 +2194,8 @@ def test_conditional_state_graph( }, "thread_id": "4", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -2261,7 +2264,8 @@ def test_conditional_state_graph( }, "thread_id": "4", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -2324,11 +2328,10 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: app = create_react_agent(model, tools) - if SHOULD_CHECK_SNAPSHOTS: - assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot - assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot + assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot + assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot assert app.invoke( {"messages": [HumanMessage(content="what is weather in sf")]} @@ -2632,9 +2635,8 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_state_graph_packets( - request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture + sync_checkpointer: BaseCheckpointSaver, mocker: MockerFixture ) -> None: from langchain_core.language_models.fake_chat_models import ( FakeMessagesListChatModel, @@ -2648,10 +2650,6 @@ def test_state_graph_packets( ) from langchain_core.tools import tool - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - class AgentState(TypedDict): messages: Annotated[list[BaseMessage], add_messages] @@ -2869,7 +2867,7 @@ def test_state_graph_packets( # interrupt after agent app_w_interrupt = workflow.compile( - checkpointer=checkpointer, + checkpointer=sync_checkpointer, interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -2947,7 +2945,8 @@ def test_state_graph_packets( }, "thread_id": "1", }, - parent_config=([*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config + parent_config=( + [*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config ), interrupts=(), ) @@ -3009,7 +3008,8 @@ def test_state_graph_packets( }, "thread_id": "1", }, - parent_config=([*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config + parent_config=( + [*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config ), interrupts=(), ) @@ -3126,7 +3126,8 @@ def test_state_graph_packets( }, "thread_id": "1", }, - parent_config=([*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config + parent_config=( + [*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config ), interrupts=(), ) @@ -3185,7 +3186,8 @@ def test_state_graph_packets( }, "thread_id": "1", }, - parent_config=([*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config + parent_config=( + [*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config ), interrupts=(), ) @@ -3193,7 +3195,7 @@ def test_state_graph_packets( # interrupt before tools app_w_interrupt = workflow.compile( - checkpointer=checkpointer, + checkpointer=sync_checkpointer, interrupt_before=["tools"], ) config = {"configurable": {"thread_id": "2"}} @@ -3272,7 +3274,8 @@ def test_state_graph_packets( }, "thread_id": "2", }, - parent_config=([*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config + parent_config=( + [*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config ), interrupts=(), ) @@ -3328,7 +3331,8 @@ def test_state_graph_packets( }, "thread_id": "2", }, - parent_config=([*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config + parent_config=( + [*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config ), interrupts=(), ) @@ -3443,7 +3447,8 @@ def test_state_graph_packets( }, "thread_id": "2", }, - parent_config=([*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config + parent_config=( + [*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config ), interrupts=(), ) @@ -3502,18 +3507,17 @@ def test_state_graph_packets( }, "thread_id": "2", }, - parent_config=([*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config + parent_config=( + [*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config ), interrupts=(), ) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_message_graph( snapshot: SnapshotAssertion, deterministic_uuids: MockerFixture, - request: pytest.FixtureRequest, - checkpointer_name: str, + sync_checkpointer: BaseCheckpointSaver, ) -> None: from copy import deepcopy @@ -3525,10 +3529,6 @@ def test_message_graph( from langchain_core.outputs import ChatGeneration, ChatResult from langchain_core.tools import tool - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - class FakeFuntionChatModel(FakeMessagesListChatModel): def bind_functions(self, functions: list): return self @@ -3634,7 +3634,7 @@ def test_message_graph( # meaning you can use it as you would any other runnable app = workflow.compile() - if SHOULD_CHECK_SNAPSHOTS and checkpointer_name == "memory": + if isinstance(sync_checkpointer, InMemorySaver): assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot @@ -3728,7 +3728,7 @@ def test_message_graph( ] app_w_interrupt = workflow.compile( - checkpointer=checkpointer, + checkpointer=sync_checkpointer, interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -3796,7 +3796,8 @@ def test_message_graph( }, "thread_id": "1", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -3845,7 +3846,8 @@ def test_message_graph( }, "thread_id": "1", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -3936,7 +3938,8 @@ def test_message_graph( }, "thread_id": "1", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -3985,13 +3988,14 @@ def test_message_graph( "writes": {"agent": AIMessage(content="answer", id="ai2")}, "thread_id": "1", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) app_w_interrupt = workflow.compile( - checkpointer=checkpointer, + checkpointer=sync_checkpointer, interrupt_before=["tools"], ) config = {"configurable": {"thread_id": "2"}} @@ -4058,7 +4062,8 @@ def test_message_graph( }, "thread_id": "2", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -4113,7 +4118,8 @@ def test_message_graph( }, "thread_id": "2", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -4204,7 +4210,8 @@ def test_message_graph( }, "thread_id": "2", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -4254,7 +4261,8 @@ def test_message_graph( "writes": {"agent": AIMessage(content="answer", id="ai2")}, "thread_id": "2", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -4304,17 +4312,16 @@ def test_message_graph( "writes": {"tools": UnsortedSequence("ai", "an extra message")}, "thread_id": "2", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_root_graph( deterministic_uuids: MockerFixture, - request: pytest.FixtureRequest, - checkpointer_name: str, + sync_checkpointer: BaseCheckpointSaver, ) -> None: from copy import deepcopy @@ -4331,10 +4338,6 @@ def test_root_graph( from langchain_core.outputs import ChatGeneration, ChatResult from langchain_core.tools import tool - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - class FakeFuntionChatModel(FakeMessagesListChatModel): def bind_functions(self, functions: list): return self @@ -4533,7 +4536,7 @@ def test_root_graph( ] app_w_interrupt = workflow.compile( - checkpointer=checkpointer, + checkpointer=sync_checkpointer, interrupt_after=["agent"], ) config = {"configurable": {"thread_id": "1"}} @@ -4601,7 +4604,8 @@ def test_root_graph( }, "thread_id": "1", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -4650,7 +4654,8 @@ def test_root_graph( }, "thread_id": "1", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -4742,7 +4747,8 @@ def test_root_graph( }, "thread_id": "1", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -4792,13 +4798,14 @@ def test_root_graph( "writes": {"agent": AIMessage(content="answer", id="ai2")}, "thread_id": "1", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) app_w_interrupt = workflow.compile( - checkpointer=checkpointer, + checkpointer=sync_checkpointer, interrupt_before=["tools"], ) config = {"configurable": {"thread_id": "2"}} @@ -4865,7 +4872,8 @@ def test_root_graph( }, "thread_id": "2", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -4920,7 +4928,8 @@ def test_root_graph( }, "thread_id": "2", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -5012,7 +5021,8 @@ def test_root_graph( }, "thread_id": "2", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -5061,7 +5071,8 @@ def test_root_graph( "writes": {"agent": AIMessage(content="answer", id="ai2")}, "thread_id": "2", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -5111,7 +5122,8 @@ def test_root_graph( "writes": {"tools": UnsortedSequence("ai", "an extra message")}, "thread_id": "2", }, - parent_config=(list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config + parent_config=( + list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config ), interrupts=(), ) @@ -5147,7 +5159,7 @@ def test_root_graph( }, ) new_workflow.add_edge("tools", "agent") - new_app = new_workflow.compile(checkpointer=checkpointer) + new_app = new_workflow.compile(checkpointer=sync_checkpointer) model.i = 0 # reset the llm # previous state is converted to new schema @@ -5192,8 +5204,7 @@ def test_root_graph( "writes": {"tools": UnsortedSequence("ai", "an extra message")}, "thread_id": "2", }, - parent_config=(list(new_app.checkpointer.list(config, limit=2))[-1].config - ), + parent_config=(list(new_app.checkpointer.list(config, limit=2))[-1].config), interrupts=(), ) @@ -5461,12 +5472,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None: ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_dynamic_interrupt( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - +def test_dynamic_interrupt(sync_checkpointer: BaseCheckpointSaver) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -5509,7 +5515,7 @@ def test_dynamic_interrupt( "market": "US", } - tool_two = tool_two_graph.compile(checkpointer=checkpointer) + tool_two = tool_two_graph.compile(checkpointer=sync_checkpointer) # missing thread_id with pytest.raises(ValueError, match="thread_id"): @@ -5596,8 +5602,7 @@ def test_dynamic_interrupt( "writes": None, "thread_id": "1", }, - parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config - ), + parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config), interrupts=( Interrupt( value="Just because...", @@ -5628,18 +5633,12 @@ def test_dynamic_interrupt( "writes": {}, "thread_id": "1", }, - parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config - ), + parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config), interrupts=(), ) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_copy_checkpoint( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - +def test_copy_checkpoint(sync_checkpointer: BaseCheckpointSaver) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -5690,7 +5689,7 @@ def test_copy_checkpoint( "market": "US", } - tool_two = tool_two_graph.compile(checkpointer=checkpointer) + tool_two = tool_two_graph.compile(checkpointer=sync_checkpointer) # missing thread_id with pytest.raises(ValueError, match="thread_id"): @@ -5786,8 +5785,7 @@ def test_copy_checkpoint( "writes": None, "thread_id": "1", }, - parent_config=([*tool_two.checkpointer.list(thread1, limit=2)][-1].config - ), + parent_config=([*tool_two.checkpointer.list(thread1, limit=2)][-1].config), interrupts=( Interrupt( value="Just because...", @@ -5797,7 +5795,6 @@ def test_copy_checkpoint( ), ) - # clear the interrupt and next tasks tool_two.update_state(thread1, None, as_node="__copy__") # interrupt is cleared, next task is kept @@ -5842,12 +5839,7 @@ def test_copy_checkpoint( ) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_dynamic_interrupt_subgraph( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - +def test_dynamic_interrupt_subgraph(sync_checkpointer: BaseCheckpointSaver) -> None: class SubgraphState(TypedDict): my_key: str market: str @@ -5902,7 +5894,7 @@ def test_dynamic_interrupt_subgraph( "market": "US", } - tool_two = tool_two_graph.compile(checkpointer=checkpointer) + tool_two = tool_two_graph.compile(checkpointer=sync_checkpointer) # missing thread_id with pytest.raises(ValueError, match="thread_id"): @@ -6004,7 +5996,8 @@ def test_dynamic_interrupt_subgraph( "writes": None, "thread_id": "1", }, - parent_config=(list( + parent_config=( + list( tool_two.checkpointer.list( {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}, limit=2 ) @@ -6040,7 +6033,8 @@ def test_dynamic_interrupt_subgraph( "writes": {}, "thread_id": "1", }, - parent_config=(list( + parent_config=( + list( tool_two.checkpointer.list( {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}, limit=2 ) @@ -6050,12 +6044,11 @@ def test_dynamic_interrupt_subgraph( ) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_start_branch_then( - snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str + snapshot: SnapshotAssertion, + sync_checkpointer: BaseCheckpointSaver, + sync_store: BaseStore, ) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -6075,7 +6068,9 @@ def test_start_branch_then( path_map=["tool_two_slow", "tool_two_fast"], ) tool_two = tool_two_graph.compile() - if checkpointer_name == "memory": + if isinstance(sync_checkpointer, InMemorySaver) and isinstance( + sync_store, InMemoryStore + ): assert tool_two.get_graph().draw_mermaid() == snapshot assert tool_two.invoke({"my_key": "value", "market": "DE"}) == { @@ -6088,8 +6083,8 @@ def test_start_branch_then( } tool_two = tool_two_graph.compile( - store=InMemoryStore(), - checkpointer=checkpointer, + store=sync_store, + checkpointer=sync_checkpointer, interrupt_before=["tool_two_fast", "tool_two_slow"], ) @@ -6143,8 +6138,7 @@ def test_start_branch_then( "assistant_id": "a", "thread_id": "1", }, - parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config - ), + parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config), interrupts=(), ) # resume, for same result as above @@ -6172,8 +6166,7 @@ def test_start_branch_then( "assistant_id": "a", "thread_id": "1", }, - parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config - ), + parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config), interrupts=(), ) @@ -6203,8 +6196,7 @@ def test_start_branch_then( "assistant_id": "a", "thread_id": "2", }, - parent_config=(list(tool_two.checkpointer.list(thread2, limit=2))[-1].config - ), + parent_config=(list(tool_two.checkpointer.list(thread2, limit=2))[-1].config), interrupts=(), ) # resume, for same result as above @@ -6232,8 +6224,7 @@ def test_start_branch_then( "assistant_id": "a", "thread_id": "2", }, - parent_config=(list(tool_two.checkpointer.list(thread2, limit=2))[-1].config - ), + parent_config=(list(tool_two.checkpointer.list(thread2, limit=2))[-1].config), interrupts=(), ) @@ -6263,8 +6254,7 @@ def test_start_branch_then( "assistant_id": "b", "thread_id": "3", }, - parent_config=(list(tool_two.checkpointer.list(thread3, limit=2))[-1].config - ), + parent_config=(list(tool_two.checkpointer.list(thread3, limit=2))[-1].config), interrupts=(), ) # update state @@ -6289,8 +6279,7 @@ def test_start_branch_then( "assistant_id": "b", "thread_id": "3", }, - parent_config=(list(tool_two.checkpointer.list(thread3, limit=2))[-1].config - ), + parent_config=(list(tool_two.checkpointer.list(thread3, limit=2))[-1].config), interrupts=(), ) # resume, for same result as above @@ -6318,18 +6307,14 @@ def test_start_branch_then( "assistant_id": "b", "thread_id": "3", }, - parent_config=(list(tool_two.checkpointer.list(thread3, limit=2))[-1].config - ), + parent_config=(list(tool_two.checkpointer.list(thread3, limit=2))[-1].config), interrupts=(), ) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_branch_then( - snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str + snapshot: SnapshotAssertion, sync_checkpointer: BaseCheckpointSaver ) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -6349,7 +6334,7 @@ def test_branch_then( tool_two_graph.add_node("finish", lambda s: {"my_key": " finished"}) tool_two = tool_two_graph.compile() - if checkpointer_name == "memory": + if isinstance(sync_checkpointer, InMemorySaver): assert tool_two.get_graph().draw_mermaid(with_styles=False) == snapshot assert tool_two.get_graph().draw_mermaid() == snapshot @@ -6363,7 +6348,7 @@ def test_branch_then( } # test stream_mode=debug - tool_two = tool_two_graph.compile(checkpointer=checkpointer) + tool_two = tool_two_graph.compile(checkpointer=sync_checkpointer) thread10 = {"configurable": {"thread_id": "10"}} res = [ @@ -6663,7 +6648,8 @@ def test_branch_then( ] tool_two = tool_two_graph.compile( - checkpointer=checkpointer, interrupt_before=["tool_two_fast", "tool_two_slow"] + checkpointer=sync_checkpointer, + interrupt_before=["tool_two_fast", "tool_two_slow"], ) # missing thread_id @@ -6695,8 +6681,7 @@ def test_branch_then( "writes": {"prepare": {"my_key": " prepared"}}, "thread_id": "1", }, - parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config - ), + parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config), interrupts=(), ) # resume, for same result as above @@ -6723,8 +6708,7 @@ def test_branch_then( "writes": {"finish": {"my_key": " finished"}}, "thread_id": "1", }, - parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config - ), + parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config), interrupts=(), ) @@ -6753,8 +6737,7 @@ def test_branch_then( "writes": {"prepare": {"my_key": " prepared"}}, "thread_id": "2", }, - parent_config=(list(tool_two.checkpointer.list(thread2, limit=2))[-1].config - ), + parent_config=(list(tool_two.checkpointer.list(thread2, limit=2))[-1].config), interrupts=(), ) # resume, for same result as above @@ -6781,13 +6764,12 @@ def test_branch_then( "writes": {"finish": {"my_key": " finished"}}, "thread_id": "2", }, - parent_config=(list(tool_two.checkpointer.list(thread2, limit=2))[-1].config - ), + parent_config=(list(tool_two.checkpointer.list(thread2, limit=2))[-1].config), interrupts=(), ) tool_two = tool_two_graph.compile( - checkpointer=checkpointer, interrupt_before=["finish"] + checkpointer=sync_checkpointer, interrupt_before=["finish"] ) thread1 = {"configurable": {"thread_id": "11"}} @@ -6819,8 +6801,7 @@ def test_branch_then( "writes": {"tool_two_slow": {"my_key": " slow"}}, "thread_id": "11", }, - parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config - ), + parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config), interrupts=(), ) @@ -6848,13 +6829,12 @@ def test_branch_then( "writes": {"tool_two_slow": {"my_key": "er"}}, "thread_id": "11", }, - parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config - ), + parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config), interrupts=(), ) tool_two = tool_two_graph.compile( - checkpointer=checkpointer, interrupt_after=["prepare"] + checkpointer=sync_checkpointer, interrupt_after=["prepare"] ) # missing thread_id @@ -6886,8 +6866,7 @@ def test_branch_then( "writes": {"prepare": {"my_key": " prepared"}}, "thread_id": "21", }, - parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config - ), + parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config), interrupts=(), ) # resume, for same result as above @@ -6914,8 +6893,7 @@ def test_branch_then( "writes": {"finish": {"my_key": " finished"}}, "thread_id": "21", }, - parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config - ), + parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config), interrupts=(), ) @@ -6944,8 +6922,7 @@ def test_branch_then( "writes": {"prepare": {"my_key": " prepared"}}, "thread_id": "22", }, - parent_config=(list(tool_two.checkpointer.list(thread2, limit=2))[-1].config - ), + parent_config=(list(tool_two.checkpointer.list(thread2, limit=2))[-1].config), interrupts=(), ) # resume, for same result as above @@ -6972,8 +6949,7 @@ def test_branch_then( "writes": {"finish": {"my_key": " finished"}}, "thread_id": "22", }, - parent_config=(list(tool_two.checkpointer.list(thread2, limit=2))[-1].config - ), + parent_config=(list(tool_two.checkpointer.list(thread2, limit=2))[-1].config), interrupts=(), ) @@ -7028,8 +7004,7 @@ def test_branch_then( "writes": {"prepare": {"my_key": " prepared"}}, "thread_id": "23", }, - parent_config=(list(tool_two.checkpointer.list(thread3, limit=2))[-1].config - ), + parent_config=(list(tool_two.checkpointer.list(thread3, limit=2))[-1].config), interrupts=(), ) # resume, for same result as above @@ -7056,8 +7031,7 @@ def test_branch_then( "writes": {"finish": {"my_key": " finished"}}, "thread_id": "23", }, - parent_config=(list(tool_two.checkpointer.list(thread3, limit=2))[-1].config - ), + parent_config=(list(tool_two.checkpointer.list(thread3, limit=2))[-1].config), interrupts=(), ) @@ -7445,12 +7419,7 @@ def test_send_dedupe_on_resume( assert history[1] == expected_history[2] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_nested_graph_state( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - +def test_nested_graph_state(sync_checkpointer: BaseCheckpointSaver) -> None: class InnerState(TypedDict): my_key: str my_other_key: str @@ -7496,7 +7465,7 @@ def test_nested_graph_state( graph.add_edge("inner", "outer_2") graph.set_finish_point("outer_2") - app = graph.compile(checkpointer=checkpointer) + app = graph.compile(checkpointer=sync_checkpointer) config = {"configurable": {"thread_id": "1"}} app.invoke({"my_key": "my value"}, config, debug=True) @@ -7528,7 +7497,8 @@ def test_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -7589,7 +7559,8 @@ def test_nested_graph_state( "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": AnyStr("inner:"), @@ -7620,7 +7591,8 @@ def test_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -7664,7 +7636,8 @@ def test_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -7776,7 +7749,8 @@ def test_nested_graph_state( "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": AnyStr("inner:"), @@ -7904,7 +7878,8 @@ def test_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -7938,7 +7913,8 @@ def test_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -8093,12 +8069,9 @@ def test_nested_graph_state( assert app.get_state(actual_snapshot.config) == expected_snapshot -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_doubly_nested_graph_state( - request: pytest.FixtureRequest, checkpointer_name: str + sync_checkpointer: BaseCheckpointSaver, ) -> None: - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - class State(TypedDict): my_key: str @@ -8146,7 +8119,7 @@ def test_doubly_nested_graph_state( graph.add_edge("child", "parent_2") graph.set_finish_point("parent_2") - app = graph.compile(checkpointer=checkpointer) + app = graph.compile(checkpointer=sync_checkpointer) # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} @@ -8191,7 +8164,8 @@ def test_doubly_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -8244,7 +8218,8 @@ def test_doubly_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": AnyStr("child:"), @@ -8303,7 +8278,8 @@ def test_doubly_nested_graph_state( "langgraph_triggers": ["branch:to:child_1"], }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": AnyStr(), @@ -8386,7 +8362,8 @@ def test_doubly_nested_graph_state( ], }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": AnyStr(), @@ -8431,7 +8408,8 @@ def test_doubly_nested_graph_state( "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": AnyStr("child:"), @@ -8462,7 +8440,8 @@ def test_doubly_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -8507,7 +8486,8 @@ def test_doubly_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -8519,7 +8499,6 @@ def test_doubly_nested_graph_state( ) ) - # get outer graph history outer_history = list(app.get_state_history(config)) assert outer_history == [ @@ -9071,12 +9050,7 @@ def test_doubly_nested_graph_state( ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_send_to_nested_graphs( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - +def test_send_to_nested_graphs(sync_checkpointer: BaseCheckpointSaver) -> None: class OverallState(TypedDict): subjects: list[str] jokes: Annotated[list[str], operator.add] @@ -9110,7 +9084,7 @@ def test_send_to_nested_graphs( builder.add_conditional_edges(START, continue_to_jokes) builder.add_edge("generate_joke", END) - graph = builder.compile(checkpointer=checkpointer) + graph = builder.compile(checkpointer=sync_checkpointer) config = {"configurable": {"thread_id": "1"}} tracer = FakeTracer() @@ -9139,14 +9113,11 @@ def test_send_to_nested_graphs( ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_send_react_interrupt( - request: pytest.FixtureRequest, checkpointer_name: str + sync_checkpointer: BaseCheckpointSaver, ) -> None: from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - ai_message = AIMessage( "", id="ai1", @@ -9200,7 +9171,7 @@ def test_send_react_interrupt( # simple interrupt-resume flow foo_called = 0 - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"]) + graph = builder.compile(checkpointer=sync_checkpointer, interrupt_before=["foo"]) thread1 = {"configurable": {"thread_id": "1"}} assert graph.invoke({"messages": [HumanMessage("hello")]}, thread1) == { "messages": [ @@ -9243,7 +9214,7 @@ def test_send_react_interrupt( # interrupt-update-resume flow foo_called = 0 - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"]) + graph = builder.compile(checkpointer=sync_checkpointer, interrupt_before=["foo"]) thread1 = {"configurable": {"thread_id": "2"}} assert graph.invoke({"messages": [HumanMessage("hello")]}, thread1) == { "messages": [ @@ -9313,7 +9284,8 @@ def test_send_react_interrupt( "thread_id": "2", }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "2", "checkpoint_ns": "", @@ -9374,7 +9346,8 @@ def test_send_react_interrupt( "thread_id": "2", }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "2", "checkpoint_ns": "", @@ -9397,7 +9370,7 @@ def test_send_react_interrupt( # interrupt-update-resume flow, creating new Send in update call foo_called = 0 - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"]) + graph = builder.compile(checkpointer=sync_checkpointer, interrupt_before=["foo"]) thread1 = {"configurable": {"thread_id": "3"}} assert graph.invoke({"messages": [HumanMessage("hello")]}, thread1) == { "messages": [ @@ -9469,7 +9442,8 @@ def test_send_react_interrupt( "thread_id": "3", }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "3", "checkpoint_ns": "", @@ -9558,7 +9532,8 @@ def test_send_react_interrupt( "thread_id": "3", }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "3", "checkpoint_ns": "", @@ -9602,14 +9577,11 @@ def test_send_react_interrupt( assert foo_called == 1 -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_send_react_interrupt_control( - request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion + sync_checkpointer: BaseCheckpointSaver, snapshot: SnapshotAssertion ) -> None: from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - ai_message = AIMessage( "", id="ai1", @@ -9635,7 +9607,7 @@ def test_send_react_interrupt_control( builder.add_edge(START, "agent") graph = builder.compile() - if checkpointer_name == "memory": + if isinstance(sync_checkpointer, InMemorySaver): assert graph.get_graph().draw_mermaid() == snapshot assert graph.invoke({"messages": [HumanMessage("hello")]}) == { @@ -9662,7 +9634,7 @@ def test_send_react_interrupt_control( # simple interrupt-resume flow foo_called = 0 - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"]) + graph = builder.compile(checkpointer=sync_checkpointer, interrupt_before=["foo"]) thread1 = {"configurable": {"thread_id": "1"}} assert graph.invoke({"messages": [HumanMessage("hello")]}, thread1) == { "messages": [ @@ -9705,7 +9677,7 @@ def test_send_react_interrupt_control( # interrupt-update-resume flow foo_called = 0 - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"]) + graph = builder.compile(checkpointer=sync_checkpointer, interrupt_before=["foo"]) thread1 = {"configurable": {"thread_id": "2"}} assert graph.invoke({"messages": [HumanMessage("hello")]}, thread1) == { "messages": [ @@ -9775,7 +9747,8 @@ def test_send_react_interrupt_control( "thread_id": "2", }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "2", "checkpoint_ns": "", @@ -9836,7 +9809,8 @@ def test_send_react_interrupt_control( "thread_id": "2", }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "2", "checkpoint_ns": "", @@ -9862,9 +9836,8 @@ def test_send_react_interrupt_control( # TODO add here test with invoke(Command()) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_weather_subgraph( - request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion + sync_checkpointer: BaseCheckpointSaver, snapshot: SnapshotAssertion ) -> None: from langchain_core.language_models.fake_chat_models import ( FakeMessagesListChatModel, @@ -9874,8 +9847,6 @@ def test_weather_subgraph( from langgraph.graph import MessagesState - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - # setup subgraph @tool @@ -9970,9 +9941,9 @@ def test_weather_subgraph( ) graph.add_edge("normal_llm_node", END) graph.add_edge("weather_graph", END) - graph = graph.compile(checkpointer=checkpointer) + graph = graph.compile(checkpointer=sync_checkpointer) - if checkpointer_name == "memory": + if isinstance(sync_checkpointer, InMemorySaver): assert graph.get_graph(xray=1).draw_mermaid() == snapshot config = {"configurable": {"thread_id": "1"}} @@ -10023,7 +9994,8 @@ def test_weather_subgraph( "thread_id": "1", }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -10113,7 +10085,8 @@ def test_weather_subgraph( "thread_id": "14", }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "14", "checkpoint_ns": "", @@ -10160,7 +10133,8 @@ def test_weather_subgraph( "langgraph_checkpoint_ns": AnyStr("weather_graph:"), }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "14", "checkpoint_ns": AnyStr("weather_graph:"), @@ -10214,7 +10188,8 @@ def test_weather_subgraph( "thread_id": "14", }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "14", "checkpoint_ns": "", @@ -10268,7 +10243,8 @@ def test_weather_subgraph( "langgraph_checkpoint_ns": AnyStr("weather_graph:"), }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "14", "checkpoint_ns": AnyStr("weather_graph:"), diff --git a/libs/langgraph/tests/test_large_cases_async.py b/libs/langgraph/tests/test_large_cases_async.py index 3a16a1034..f0ae78813 100644 --- a/libs/langgraph/tests/test_large_cases_async.py +++ b/libs/langgraph/tests/test_large_cases_async.py @@ -14,7 +14,6 @@ import pytest from langchain_core.messages import ToolCall from langchain_core.runnables import RunnableConfig, RunnablePick from pytest_mock import MockerFixture -from syrupy import SnapshotAssertion from typing_extensions import TypedDict from langgraph.channels.last_value import LastValue @@ -27,14 +26,10 @@ from langgraph.graph.state import StateGraph from langgraph.prebuilt.chat_agent_executor import create_react_agent from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import Channel, Pregel -from langgraph.store.memory import InMemoryStore +from langgraph.store.base import BaseStore from langgraph.types import PregelTask, Send, StateSnapshot, StreamWriter from tests.any_int import AnyInt from tests.any_str import AnyDict, AnyStr, UnsortedSequence -from tests.conftest import ( - ALL_CHECKPOINTERS_ASYNC, - awith_checkpointer, -) from tests.fake_chat import FakeChatModel from tests.fake_tracer import FakeTracer from tests.messages import ( @@ -47,274 +42,262 @@ from tests.messages import ( pytestmark = pytest.mark.anyio -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_invoke_two_processes_in_out_interrupt( - checkpointer_name: str, mocker: MockerFixture + async_checkpointer: BaseCheckpointSaver, mocker: MockerFixture ) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") - async with awith_checkpointer(checkpointer_name) as checkpointer: - app = Pregel( - nodes={"one": one, "two": two}, - channels={ - "inbox": LastValue(int), - "output": LastValue(int), - "input": LastValue(int), + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "inbox": LastValue(int), + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", + checkpointer=async_checkpointer, + interrupt_after_nodes=["one"], + ) + thread1 = {"configurable": {"thread_id": "1"}} + thread2 = {"configurable": {"thread_id": "2"}} + + # start execution, stop at inbox + assert await app.ainvoke(2, thread1) is None + + # inbox == 3 + checkpoint = await async_checkpointer.aget(thread1) + assert checkpoint is not None + assert checkpoint["channel_values"]["inbox"] == 3 + + # resume execution, finish + assert await app.ainvoke(None, thread1) == 4 + + # start execution again, stop at inbox + assert await app.ainvoke(20, thread1) is None + + # inbox == 21 + checkpoint = await async_checkpointer.aget(thread1) + assert checkpoint is not None + assert checkpoint["channel_values"]["inbox"] == 21 + + # send a new value in, interrupting the previous execution + assert await app.ainvoke(3, thread1) is None + assert await app.ainvoke(None, thread1) == 5 + + # start execution again, stopping at inbox + assert await app.ainvoke(20, thread2) is None + + # inbox == 21 + snapshot = await app.aget_state(thread2) + assert snapshot.values["inbox"] == 21 + assert snapshot.next == ("two",) + + # update the state, resume + await app.aupdate_state(thread2, 25, as_node="one") + assert await app.ainvoke(None, thread2) == 26 + + # no pending tasks + snapshot = await app.aget_state(thread2) + assert snapshot.next == () + + # list history + history = [c async for c in app.aget_state_history(thread1)] + assert history == [ + StateSnapshot( + values={"inbox": 4, "output": 5, "input": 3}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, - input_channels="input", - output_channels="output", - checkpointer=checkpointer, - interrupt_after_nodes=["one"], - ) - thread1 = {"configurable": {"thread_id": "1"}} - thread2 = {"configurable": {"thread_id": "2"}} - - # start execution, stop at inbox - assert await app.ainvoke(2, thread1) is None - - # inbox == 3 - checkpoint = await checkpointer.aget(thread1) - assert checkpoint is not None - assert checkpoint["channel_values"]["inbox"] == 3 - - # resume execution, finish - assert await app.ainvoke(None, thread1) == 4 - - # start execution again, stop at inbox - assert await app.ainvoke(20, thread1) is None - - # inbox == 21 - checkpoint = await checkpointer.aget(thread1) - assert checkpoint is not None - assert checkpoint["channel_values"]["inbox"] == 21 - - # send a new value in, interrupting the previous execution - assert await app.ainvoke(3, thread1) is None - assert await app.ainvoke(None, thread1) == 5 - - # start execution again, stopping at inbox - assert await app.ainvoke(20, thread2) is None - - # inbox == 21 - snapshot = await app.aget_state(thread2) - assert snapshot.values["inbox"] == 21 - assert snapshot.next == ("two",) - - # update the state, resume - await app.aupdate_state(thread2, 25, as_node="one") - assert await app.ainvoke(None, thread2) == 26 - - # no pending tasks - snapshot = await app.aget_state(thread2) - assert snapshot.next == () - - # list history - history = [c async for c in app.aget_state_history(thread1)] - assert history == [ - StateSnapshot( - values={"inbox": 4, "output": 5, "input": 3}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 6, - "writes": {"two": 5}, + metadata={ + "parents": {}, + "source": "loop", + "step": 6, + "writes": {"two": 5}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[1].config, + interrupts=(), + ), + StateSnapshot( + values={"inbox": 4, "output": 4, "input": 3}, + tasks=(PregelTask(AnyStr(), "two", (PULL, "two"), result={"output": 5}),), + next=("two",), + config={ + "configurable": { "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[1].config, - interrupts=(), - ), - StateSnapshot( - values={"inbox": 4, "output": 4, "input": 3}, - tasks=( - PregelTask(AnyStr(), "two", (PULL, "two"), result={"output": 5}), - ), - next=("two",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 5, - "writes": {"one": None}, + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 5, + "writes": {"one": None}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[2].config, + interrupts=(), + ), + StateSnapshot( + values={"inbox": 21, "output": 4, "input": 3}, + tasks=(PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 4}),), + next=("one",), + config={ + "configurable": { "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[2].config, - interrupts=(), - ), - StateSnapshot( - values={"inbox": 21, "output": 4, "input": 3}, - tasks=( - PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 4}), - ), - next=("one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "step": 4, - "writes": {"input": 3}, + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "step": 4, + "writes": {"input": 3}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[3].config, + interrupts=(), + ), + StateSnapshot( + values={"inbox": 21, "output": 4, "input": 20}, + tasks=(PregelTask(AnyStr(), "two", (PULL, "two")),), + next=("two",), + config={ + "configurable": { "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[3].config, - interrupts=(), - ), - StateSnapshot( - values={"inbox": 21, "output": 4, "input": 20}, - tasks=(PregelTask(AnyStr(), "two", (PULL, "two")),), - next=("two",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"one": None}, + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"one": None}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[4].config, + interrupts=(), + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 20}, + tasks=(PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 21}),), + next=("one",), + config={ + "configurable": { "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[4].config, - interrupts=(), - ), - StateSnapshot( - values={"inbox": 3, "output": 4, "input": 20}, - tasks=( - PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 21}), - ), - next=("one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "step": 2, - "writes": {"input": 20}, + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "step": 2, + "writes": {"input": 20}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[5].config, + interrupts=(), + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 2}, + tasks=(), + next=(), + config={ + "configurable": { "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[5].config, - interrupts=(), - ), - StateSnapshot( - values={"inbox": 3, "output": 4, "input": 2}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"two": 4}, + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"two": 4}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[6].config, + interrupts=(), + ), + StateSnapshot( + values={"inbox": 3, "input": 2}, + tasks=(PregelTask(AnyStr(), "two", (PULL, "two"), result={"output": 4}),), + next=("two",), + config={ + "configurable": { "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[6].config, - interrupts=(), - ), - StateSnapshot( - values={"inbox": 3, "input": 2}, - tasks=( - PregelTask(AnyStr(), "two", (PULL, "two"), result={"output": 4}), - ), - next=("two",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": {"one": None}, + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": {"one": None}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=history[7].config, + interrupts=(), + ), + StateSnapshot( + values={"input": 2}, + tasks=(PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 3}),), + next=("one",), + config={ + "configurable": { "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=history[7].config, - interrupts=(), - ), - StateSnapshot( - values={"input": 2}, - tasks=( - PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 3}), - ), - next=("one",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "step": -1, - "writes": {"input": 2}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - interrupts=(), - ), - ] + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "step": -1, + "writes": {"input": 2}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ), + ] - # forking from any previous checkpoint should re-run nodes - assert [ - c async for c in app.astream(None, history[0].config, stream_mode="updates") - ] == [] - assert [ - c async for c in app.astream(None, history[1].config, stream_mode="updates") - ] == [ - {"two": {"output": 5}}, - ] - assert [ - c async for c in app.astream(None, history[2].config, stream_mode="updates") - ] == [ - {"one": {"inbox": 4}}, - {"__interrupt__": ()}, - ] + # forking from any previous checkpoint should re-run nodes + assert [ + c async for c in app.astream(None, history[0].config, stream_mode="updates") + ] == [] + assert [ + c async for c in app.astream(None, history[1].config, stream_mode="updates") + ] == [ + {"two": {"output": 5}}, + ] + assert [ + c async for c in app.astream(None, history[2].config, stream_mode="updates") + ] == [ + {"one": {"inbox": 4}}, + {"__interrupt__": ()}, + ] async def test_fork_always_re_runs_nodes( @@ -523,8 +506,7 @@ async def test_fork_always_re_runs_nodes( ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_conditional_graph(checkpointer_name: str) -> None: +async def test_conditional_graph(async_checkpointer: BaseCheckpointSaver) -> None: from langchain_core.agents import AgentAction, AgentFinish from langchain_core.language_models.fake import FakeStreamingListLLM from langchain_core.prompts import PromptTemplate @@ -787,77 +769,88 @@ async def test_conditional_graph(checkpointer_name: str) -> None: }, ] - async with awith_checkpointer(checkpointer_name) as checkpointer: - # test state get/update methods with interrupt_after + # test state get/update methods with interrupt_after - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["agent"], + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_after=["agent"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + {"input": "what is weather in sf"}, config ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c - async for c in app_w_interrupt.astream( - {"input": "what is weather in sf"}, config - ) - ] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } + ] == [ + { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), } - ] + } + ] - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + }, + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, + created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ + "ts" + ], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": { "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - }, - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, - created_at=( - await app_w_interrupt.checkpointer.aget_tuple(config) - ).checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": { "agent": { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), } - }, - "thread_id": "1", + } }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - interrupts=(), - ) + "thread_id": "1", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + interrupts=(), + ) - await app_w_interrupt.aupdate_state( - config, - { + await app_w_interrupt.aupdate_state( + config, + { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + }, + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent": { "agent_outcome": AgentAction( tool="search_api", tool_input="query", @@ -865,10 +858,22 @@ async def test_conditional_graph(checkpointer_name: str) -> None: ), "input": "what is weather in sf", }, - ) - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": { "agent": { "agent_outcome": AgentAction( tool="search_api", @@ -876,91 +881,88 @@ async def test_conditional_graph(checkpointer_name: str) -> None: log="tool:search_api:a different query", ), "input": "what is weather in sf", - }, - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), } }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": { - "agent": { - "agent_outcome": AgentAction( + "thread_id": "1", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + interrupts=(), + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + }, + }, + { + "tools": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), - "input": "what is weather in sf", - } - }, - "thread_id": "1", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - interrupts=(), - ) + "result for query", + ] + ], + } + }, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + ] - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "agent": { - "agent_outcome": AgentAction( + await app_w_interrupt.aupdate_state( + config, + { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), - "input": "what is weather in sf", - }, - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - ] + "result for query", + ] + ], + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + }, + ) - await app_w_interrupt.aupdate_state( - config, - { + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent": { "input": "what is weather in sf", "intermediate_steps": [ [ @@ -977,10 +979,22 @@ async def test_conditional_graph(checkpointer_name: str) -> None: log="finish:a really nice answer", ), }, - ) - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ + }, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 4, + "writes": { "agent": { "input": "what is weather in sf", "intermediate_steps": [ @@ -997,124 +1011,103 @@ async def test_conditional_graph(checkpointer_name: str) -> None: return_values={"answer": "a really nice answer"}, log="finish:a really nice answer", ), - }, - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), } }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 4, - "writes": { + "thread_id": "1", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + interrupts=(), + ) + + # test state get/update methods with interrupt_before + + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_before=["tools"], + ) + config = {"configurable": {"thread_id": "2"}} + llm.i = 0 + + assert [ + c + async for c in app_w_interrupt.astream( + {"input": "what is weather in sf"}, config + ) + ] == [ + { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + } + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + }, + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": { + "agent": { "agent": { "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", ), } - }, - "thread_id": "1", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - interrupts=(), - ) - - # test state get/update methods with interrupt_before - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_before=["tools"], - ) - config = {"configurable": {"thread_id": "2"}} - llm.i = 0 - - assert [ - c - async for c in app_w_interrupt.astream( - {"input": "what is weather in sf"}, config - ) - ] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - } - ] - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - }, - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), } }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": { - "agent": { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - } - }, - "thread_id": "2", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - interrupts=(), - ) + "thread_id": "2", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + interrupts=(), + ) - await app_w_interrupt.aupdate_state( - config, - { + await app_w_interrupt.aupdate_state( + config, + { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + }, + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent": { "agent_outcome": AgentAction( tool="search_api", tool_input="query", @@ -1122,10 +1115,22 @@ async def test_conditional_graph(checkpointer_name: str) -> None: ), "input": "what is weather in sf", }, - ) - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": { "agent": { "agent_outcome": AgentAction( tool="search_api", @@ -1133,91 +1138,88 @@ async def test_conditional_graph(checkpointer_name: str) -> None: log="tool:search_api:a different query", ), "input": "what is weather in sf", - }, - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), } }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": { - "agent": { - "agent_outcome": AgentAction( + "thread_id": "2", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + interrupts=(), + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + }, + }, + { + "tools": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), - "input": "what is weather in sf", - } - }, - "thread_id": "2", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - interrupts=(), - ) + "result for query", + ] + ], + } + }, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + ] - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "agent": { - "agent_outcome": AgentAction( + await app_w_interrupt.aupdate_state( + config, + { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( tool="search_api", tool_input="query", log="tool:search_api:a different query", ), - "input": "what is weather in sf", - }, - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - ] + "result for query", + ] + ], + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + }, + ) - await app_w_interrupt.aupdate_state( - config, - { + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent": { "input": "what is weather in sf", "intermediate_steps": [ [ @@ -1234,10 +1236,22 @@ async def test_conditional_graph(checkpointer_name: str) -> None: log="finish:a really nice answer", ), }, - ) - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ + }, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 4, + "writes": { "agent": { "input": "what is weather in sf", "intermediate_steps": [ @@ -1254,244 +1268,210 @@ async def test_conditional_graph(checkpointer_name: str) -> None: return_values={"answer": "a really nice answer"}, log="finish:a really nice answer", ), - }, - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), } }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 4, - "writes": { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - } - }, - "thread_id": "2", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - interrupts=(), + "thread_id": "2", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + interrupts=(), + ) + + # test re-invoke to continue with interrupt_before + + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_before=["tools"], + ) + config = {"configurable": {"thread_id": "3"}} + llm.i = 0 # reset the llm + + assert [ + c + async for c in app_w_interrupt.astream( + {"input": "what is weather in sf"}, config ) - - # test re-invoke to continue with interrupt_before - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_before=["tools"], - ) - config = {"configurable": {"thread_id": "3"}} - llm.i = 0 # reset the llm - - assert [ - c - async for c in app_w_interrupt.astream( - {"input": "what is weather in sf"}, config - ) - ] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } + ] == [ + { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), } - ] + } + ] - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - }, + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "3", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": { - "agent": { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - } - }, + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { "thread_id": "3", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": { + "agent": { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + } }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - interrupts=(), - ) + "thread_id": "3", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + interrupts=(), + ) - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "agent": { - "input": "what is weather in sf", - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - }, + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] - ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - ] + }, + { + "tools": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + ] - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ] + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", + ] + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + { + "tools": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", ], - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - { - "tools": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ], - [ - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ], + [ + AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + "result for another", ], - } - }, - { - "agent": { - "input": "what is weather in sf", - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ], - [ - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ], + ], + } + }, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + "result for query", ], - "agent_outcome": AgentFinish( - return_values={"answer": "answer"}, log="finish:answer" - ), - } - }, - ] + [ + AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + "result for another", + ], + ], + "agent_outcome": AgentFinish( + return_values={"answer": "answer"}, log="finish:answer" + ), + } + }, + ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_conditional_graph_state(checkpointer_name: str) -> None: +async def test_conditional_graph_state(async_checkpointer: BaseCheckpointSaver) -> None: from langchain_core.agents import AgentAction, AgentFinish from langchain_core.language_models.fake import FakeStreamingListLLM from langchain_core.prompts import PromptTemplate @@ -1685,226 +1665,56 @@ async def test_conditional_graph_state(checkpointer_name: str) -> None: }, ] - async with awith_checkpointer(checkpointer_name) as checkpointer: - # test state get/update methods with interrupt_after + # test state get/update methods with interrupt_after - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["agent"], + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_after=["agent"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + {"input": "what is weather in sf"}, config ) - config = {"configurable": {"thread_id": "1"}} - - assert [ - c - async for c in app_w_interrupt.astream( - {"input": "what is weather in sf"}, config - ) - ] == [ - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - }, - {"__interrupt__": ()}, - ] - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ + ] == [ + { + "agent": { "agent_outcome": AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query", ), - "intermediate_steps": [], - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - }, - "thread_id": "1", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - interrupts=(), - ) + } + }, + {"__interrupt__": ()}, + ] - await app_w_interrupt.aupdate_state( - config, - { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ) - }, - ) - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "intermediate_steps": [], - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ) - } - }, - "thread_id": "1", - }, - parent_config=([ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", ), - interrupts=(), - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "tools": { - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - } - }, - { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - } - }, - {"__interrupt__": ()}, - ] - - await app_w_interrupt.aupdate_state( - config, - { - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ) - }, - ) - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": { - "agent": { - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ) - } - }, + "intermediate_steps": [], + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { "thread_id": "1", - }, - parent_config=([ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - interrupts=(), - ) - - # test state get/update methods with interrupt_before - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_before=["tools"], - ) - config = {"configurable": {"thread_id": "2"}} - llm.i = 0 # reset the llm - - assert [ - c - async for c in app_w_interrupt.astream( - {"input": "what is weather in sf"}, config - ) - ] == [ - { + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": { "agent": { "agent_outcome": AgentAction( tool="search_api", @@ -1913,142 +1723,70 @@ async def test_conditional_graph_state(checkpointer_name: str) -> None: ), } }, - {"__interrupt__": ()}, - ] + "thread_id": "1", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + interrupts=(), + ) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent_outcome": AgentAction( - tool="search_api", tool_input="query", log="tool:search_api:query" - ), - "intermediate_steps": [], - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - } - }, - "thread_id": "2", - }, - parent_config=([ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config + await app_w_interrupt.aupdate_state( + config, + { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ) + }, + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", ), - interrupts=(), - ) - - await app_w_interrupt.aupdate_state( - config, - { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ) - }, - ) - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "intermediate_steps": [], - }, - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": { - "agent_outcome": AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ) - } - }, - "thread_id": "2", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - interrupts=(), - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "tools": { - "intermediate_steps": [ - [ - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:a different query", - ), - "result for query", - ] - ], - } - }, - { + "intermediate_steps": [], + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 2, + "writes": { "agent": { "agent_outcome": AgentAction( tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), + tool_input="query", + log="tool:search_api:a different query", + ) } }, - {"__interrupt__": ()}, - ] + "thread_id": "1", + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) - await app_w_interrupt.aupdate_state( - config, - { - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ) - }, - ) - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ), + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": { "intermediate_steps": [ [ AgentAction( @@ -2059,36 +1797,280 @@ async def test_conditional_graph_state(checkpointer_name: str) -> None: "result for query", ] ], - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), + } + }, + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + {"__interrupt__": ()}, + ] + + await app_w_interrupt.aupdate_state( + config, + { + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ) + }, + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + }, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "writes": { + "agent": { + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ) } }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": { - "agent": { - "agent_outcome": AgentFinish( - return_values={"answer": "a really nice answer"}, - log="finish:a really nice answer", - ) - } - }, - "thread_id": "2", - }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, - interrupts=(), + "thread_id": "1", + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) + + # test state get/update methods with interrupt_before + + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_before=["tools"], + ) + config = {"configurable": {"thread_id": "2"}} + llm.i = 0 # reset the llm + + assert [ + c + async for c in app_w_interrupt.astream( + {"input": "what is weather in sf"}, config ) + ] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + }, + {"__interrupt__": ()}, + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + "intermediate_steps": [], + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:query", + ), + } + }, + "thread_id": "2", + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) + + await app_w_interrupt.aupdate_state( + config, + { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ) + }, + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "intermediate_steps": [], + }, + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 2, + "writes": { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ) + } + }, + "thread_id": "2", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + interrupts=(), + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": { + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + } + }, + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + {"__interrupt__": ()}, + ] + + await app_w_interrupt.aupdate_state( + config, + { + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ) + }, + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + "intermediate_steps": [ + [ + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ] + ], + }, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "writes": { + "agent": { + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ) + } + }, + "thread_id": "2", + }, + parent_config=[ + c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) + ][-1].config, + interrupts=(), + ) async def test_prebuilt_tool_chat() -> None: @@ -2420,8 +2402,7 @@ async def test_prebuilt_tool_chat() -> None: ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_state_graph_packets(checkpointer_name: str) -> None: +async def test_state_graph_packets(async_checkpointer: BaseCheckpointSaver) -> None: from langchain_core.language_models.fake_chat_models import ( FakeMessagesListChatModel, ) @@ -2639,107 +2620,130 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: {"agent": {"messages": AIMessage(content="answer", id="ai3")}}, ] - async with awith_checkpointer(checkpointer_name) as checkpointer: - # interrupt after agent + # interrupt after agent - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["agent"], + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_after=["agent"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + {"messages": HumanMessage(content="what is weather in sf")}, config ) - config = {"configurable": {"thread_id": "1"}} + ] == [ + { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + } + }, + {"__interrupt__": ()}, + ] - assert [ - c - async for c in app_w_interrupt.astream( - {"messages": HumanMessage(content="what is weather in sf")}, config - ) - ] == [ - { + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + ] + }, + tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0, False)),), + next=("tools",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": { "agent": { "messages": AIMessage( - id="ai1", content="", + id="ai1", tool_calls=[ { - "id": "tool_call123", "name": "search_api", "args": {"query": "query"}, - }, + "id": "tool_call123", + "type": "tool_call", + } ], ) } }, - {"__interrupt__": ()}, - ] + "thread_id": "1", + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ), - ] - }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0, False)),), - next=("tools",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": { - "messages": AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "name": "search_api", - "args": {"query": "query"}, - "id": "tool_call123", - "type": "tool_call", - } - ], - ) - } - }, - "thread_id": "1", - }, - parent_config=([ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - interrupts=(), - ) + # modify ai message + last_message = (await app_w_interrupt.aget_state(config)).values["messages"][-1] + last_message.tool_calls[0]["args"]["query"] = "a different query" + await app_w_interrupt.aupdate_state(config, {"messages": last_message}) - # modify ai message - last_message = (await app_w_interrupt.aget_state(config)).values["messages"][-1] - last_message.tool_calls[0]["args"]["query"] = "a different query" - await app_w_interrupt.aupdate_state(config, {"messages": last_message}) - - # message was replaced instead of appended - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( + # message was replaced instead of appended + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + ] + }, + tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0, False)),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 2, + "writes": { + "agent": { + "messages": AIMessage( id="ai1", content="", tool_calls=[ @@ -2749,52 +2753,103 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: "args": {"query": "a different query"}, }, ], - ), - ] - }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0, False)),), - next=("tools",), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": { - "messages": AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - }, - ], - ) - } - }, - "thread_id": "1", - }, - parent_config=([ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - interrupts=(), - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "tools": { - "messages": _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", ) } }, - { + "thread_id": "1", + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": { + "messages": _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ) + } + }, + { + "agent": { + "messages": AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], + ) + }, + }, + {"__interrupt__": ()}, + ] + + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], + ), + ] + }, + tasks=( + PregelTask(AnyStr(), "tools", (PUSH, 0, False)), + PregelTask(AnyStr(), "tools", (PUSH, 1, False)), + ), + next=("tools", "tools"), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + "writes": { "agent": { "messages": AIMessage( id="ai2", @@ -2811,240 +2866,193 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: "args": {"query": "a third one", "idx": 1}, }, ], + ), + }, + }, + "thread_id": "1", + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) + + await app_w_interrupt.aupdate_state( + config, + {"messages": AIMessage(content="answer", id="ai2")}, + ) + + # replaces message even if object identity is different, as long as id is the same + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage(content="answer", id="ai2"), + ] + }, + tasks=(), + next=(), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "writes": { + "agent": { + "messages": AIMessage(content="answer", id="ai2"), + } + }, + "thread_id": "1", + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) + + # interrupt before tools + + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_before=["tools"], + ) + config = {"configurable": {"thread_id": "2"}} + model.i = 0 + + assert [ + c + async for c in app_w_interrupt.astream( + {"messages": HumanMessage(content="what is weather in sf")}, config + ) + ] == [ + { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + } + }, + {"__interrupt__": ()}, + ] + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + ] + }, + tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0, False)),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": { + "agent": { + "messages": AIMessage( + content="", + additional_kwargs={}, + response_metadata={}, + id="ai1", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "query"}, + "id": "tool_call123", + "type": "tool_call", + } + ], ) - }, + } }, - {"__interrupt__": ()}, - ] + "thread_id": "2", + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - }, - ], - ), - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ), - AIMessage( - id="ai2", - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another", "idx": 0}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - }, - ], - ), - ] - }, - tasks=( - PregelTask(AnyStr(), "tools", (PUSH, 0, False)), - PregelTask(AnyStr(), "tools", (PUSH, 1, False)), - ), - next=("tools", "tools"), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 4, - "writes": { - "agent": { - "messages": AIMessage( - id="ai2", - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another", "idx": 0}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - }, - ], - ), - }, - }, - "thread_id": "1", - }, - parent_config=([ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - interrupts=(), - ) + # modify ai message + last_message = (await app_w_interrupt.aget_state(config)).values["messages"][-1] + last_message.tool_calls[0]["args"]["query"] = "a different query" + await app_w_interrupt.aupdate_state(config, {"messages": last_message}) - await app_w_interrupt.aupdate_state( - config, - {"messages": AIMessage(content="answer", id="ai2")}, - ) - - # replaces message even if object identity is different, as long as id is the same - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - }, - ], - ), - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ), - AIMessage(content="answer", id="ai2"), - ] - }, - tasks=(), - next=(), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": { - "agent": { - "messages": AIMessage(content="answer", id="ai2"), - } - }, - "thread_id": "1", - }, - parent_config=([ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - interrupts=(), - ) - - # interrupt before tools - - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_before=["tools"], - ) - config = {"configurable": {"thread_id": "2"}} - model.i = 0 - - assert [ - c - async for c in app_w_interrupt.astream( - {"messages": HumanMessage(content="what is weather in sf")}, config - ) - ] == [ - { + # message was replaced instead of appended + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + ] + }, + tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0, False)),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 2, + "writes": { "agent": { "messages": AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ) - } - }, - {"__interrupt__": ()}, - ] - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ), - ] - }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0, False)),), - next=("tools",), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": { - "messages": AIMessage( - content="", - additional_kwargs={}, - response_metadata={}, - id="ai1", - tool_calls=[ - { - "name": "search_api", - "args": {"query": "query"}, - "id": "tool_call123", - "type": "tool_call", - } - ], - ) - } - }, - "thread_id": "2", - }, - parent_config=([ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - interrupts=(), - ) - - # modify ai message - last_message = (await app_w_interrupt.aget_state(config)).values["messages"][-1] - last_message.tool_calls[0]["args"]["query"] = "a different query" - await app_w_interrupt.aupdate_state(config, {"messages": last_message}) - - # message was replaced instead of appended - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( id="ai1", content="", tool_calls=[ @@ -3054,52 +3062,103 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: "args": {"query": "a different query"}, }, ], - ), - ] - }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0, False)),), - next=("tools",), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": { - "messages": AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - }, - ], - ) - } - }, - "thread_id": "2", - }, - parent_config=([ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - interrupts=(), - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "tools": { - "messages": _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", ) } }, - { + "thread_id": "2", + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": { + "messages": _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ) + } + }, + { + "agent": { + "messages": AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], + ) + }, + }, + {"__interrupt__": ()}, + ] + + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], + ), + ] + }, + tasks=( + PregelTask(AnyStr(), "tools", (PUSH, 0, False)), + PregelTask(AnyStr(), "tools", (PUSH, 1, False)), + ), + next=("tools", "tools"), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + "writes": { "agent": { "messages": AIMessage( id="ai2", @@ -3116,146 +3175,74 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: "args": {"query": "a third one", "idx": 1}, }, ], - ) + ), }, }, - {"__interrupt__": ()}, - ] + "thread_id": "2", + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - }, - ], - ), - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ), - AIMessage( - id="ai2", - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another", "idx": 0}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - }, - ], - ), - ] + await app_w_interrupt.aupdate_state( + config, + {"messages": AIMessage(content="answer", id="ai2")}, + ) + + # replaces message even if object identity is different, as long as id is the same + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage(content="answer", id="ai2"), + ] + }, + tasks=(), + next=(), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "writes": { + "agent": { + "messages": AIMessage(content="answer", id="ai2"), + } }, - tasks=( - PregelTask(AnyStr(), "tools", (PUSH, 0, False)), - PregelTask(AnyStr(), "tools", (PUSH, 1, False)), - ), - next=("tools", "tools"), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 4, - "writes": { - "agent": { - "messages": AIMessage( - id="ai2", - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another", "idx": 0}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - }, - ], - ), - }, - }, - "thread_id": "2", - }, - parent_config=([ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - interrupts=(), - ) - - await app_w_interrupt.aupdate_state( - config, - {"messages": AIMessage(content="answer", id="ai2")}, - ) - - # replaces message even if object identity is different, as long as id is the same - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - }, - ], - ), - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ), - AIMessage(content="answer", id="ai2"), - ] - }, - tasks=(), - next=(), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": { - "agent": { - "messages": AIMessage(content="answer", id="ai2"), - } - }, - "thread_id": "2", - }, - parent_config=([ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - interrupts=(), - ) + "thread_id": "2", + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_message_graph(checkpointer_name: str) -> None: +async def test_message_graph(async_checkpointer: BaseCheckpointSaver) -> None: from langchain_core.language_models.fake_chat_models import ( FakeMessagesListChatModel, ) @@ -3441,20 +3428,59 @@ async def test_message_graph(checkpointer_name: str) -> None: {"agent": AIMessage(content="answer", id="ai3")}, ] - async with awith_checkpointer(checkpointer_name) as checkpointer: - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["agent"], - ) - config = {"configurable": {"thread_id": "1"}} + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_after=["agent"], + ) + config = {"configurable": {"thread_id": "1"}} - assert [ - c - async for c in app_w_interrupt.astream( - HumanMessage(content="what is weather in sf"), config + assert [ + c + async for c in app_w_interrupt.astream( + HumanMessage(content="what is weather in sf"), config + ) + ] == [ + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", ) - ] == [ - { + }, + {"__interrupt__": ()}, + ] + + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + } + ], + id="ai1", + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": { "agent": AIMessage( content="", tool_calls=[ @@ -3467,140 +3493,49 @@ async def test_message_graph(checkpointer_name: str) -> None: id="ai1", ) }, - {"__interrupt__": ()}, - ] + "thread_id": "1", + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ), - ], - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - } - ], - id="ai1", - ) - }, - "thread_id": "1", - }, - parent_config=([ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config + # modify ai message + last_message = (await app_w_interrupt.aget_state(config)).values[-1] + last_message.tool_calls[0]["args"] = {"query": "a different query"} + await app_w_interrupt.aupdate_state(config, last_message) + + # message was replaced instead of appended + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], ), - interrupts=(), - ) - - # modify ai message - last_message = (await app_w_interrupt.aget_state(config)).values[-1] - last_message.tool_calls[0]["args"] = {"query": "a different query"} - await app_w_interrupt.aupdate_state(config, last_message) - - # message was replaced instead of appended - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - ), - ], - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 2, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - id="ai1", - ) - }, - "thread_id": "1", - }, - parent_config=([ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - interrupts=(), - ) - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "tools": [ - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ) - ] - }, - { + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 2, + "writes": { "agent": AIMessage( content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ) - }, - {"__interrupt__": ()}, - ] - - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - id="ai1", tool_calls=[ { "id": "tool_call123", @@ -3608,13 +3543,87 @@ async def test_message_graph(checkpointer_name: str) -> None: "args": {"query": "a different query"}, } ], - ), + id="ai1", + ) + }, + "thread_id": "1", + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": [ _AnyIdToolMessage( content="result for a different query", name="search_api", tool_call_id="tool_call123", - ), - AIMessage( + ) + ] + }, + { + "agent": AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ) + }, + {"__interrupt__": ()}, + ] + + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call456", + "name": "search_api", + "args": {"query": "another"}, + } + ], + id="ai2", + ), + ], + tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + "writes": { + "agent": AIMessage( content="", tool_calls=[ { @@ -3624,83 +3633,64 @@ async def test_message_graph(checkpointer_name: str) -> None: } ], id="ai2", - ), - ], - tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),), - next=("tools",), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 4, - "writes": { - "agent": AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call456", - "name": "search_api", - "args": {"query": "another"}, - } - ], - id="ai2", - ) - }, - "thread_id": "1", + ) }, - parent_config=([ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - interrupts=(), - ) + "thread_id": "1", + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) - await app_w_interrupt.aupdate_state( - config, + await app_w_interrupt.aupdate_state( + config, + AIMessage(content="answer", id="ai2"), + ) + + # replaces message even if object identity is different, as long as id is the same + tup = await app_w_interrupt.checkpointer.aget_tuple(config) + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + _AnyIdHumanMessage(content="what is weather in sf"), + AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + } + ], + ), + _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), AIMessage(content="answer", id="ai2"), - ) - - # replaces message even if object identity is different, as long as id is the same - tup = await app_w_interrupt.checkpointer.aget_tuple(config) - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values=[ - _AnyIdHumanMessage(content="what is weather in sf"), - AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "a different query"}, - } - ], - ), - _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ), - AIMessage(content="answer", id="ai2"), - ], - tasks=(), - next=(), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 5, - "writes": {"agent": AIMessage(content="answer", id="ai2")}, - "thread_id": "1", - }, - parent_config=([ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config - ), - interrupts=(), - ) + ], + tasks=(), + next=(), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 5, + "writes": {"agent": AIMessage(content="answer", id="ai2")}, + "thread_id": "1", + }, + parent_config=( + [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) async def test_in_one_fan_out_out_one_graph_state() -> None: @@ -3929,8 +3919,9 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_start_branch_then(checkpointer_name: str) -> None: +async def test_start_branch_then( + async_checkpointer: BaseCheckpointSaver, async_store: BaseStore +) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -3958,257 +3949,250 @@ async def test_start_branch_then(checkpointer_name: str) -> None: "market": "US", } - async with awith_checkpointer(checkpointer_name) as checkpointer: - tool_two = tool_two_graph.compile( - store=InMemoryStore(), - checkpointer=checkpointer, - interrupt_before=["tool_two_fast", "tool_two_slow"], - ) + tool_two = tool_two_graph.compile( + store=async_store, + checkpointer=async_checkpointer, + interrupt_before=["tool_two_fast", "tool_two_slow"], + ) - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - thread1 = {"configurable": {"thread_id": "1", "assistant_id": "a"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == { - "my_key": "value", - "market": "DE", - } - assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "assistant_id": "a", - "thread_id": "1", - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value", "market": "DE"}}, - "assistant_id": "a", - "thread_id": "1", - }, - ] + thread1 = {"configurable": {"thread_id": "1", "assistant_id": "a"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == { + "my_key": "value", + "market": "DE", + } + assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "assistant_id": "a", + "thread_id": "1", + }, + { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value", "market": "DE"}}, + "assistant_id": "a", + "thread_id": "1", + }, + ] - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), - next=("tool_two_slow",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "assistant_id": "a", + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), + next=("tool_two_slow",), + config={ + "configurable": { "thread_id": "1", - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread1, limit=2) - ][-1].config, - interrupts=(), - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread1, debug=1) == { - "my_key": "value slow", - "market": "DE", - } - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value slow", "market": "DE"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"tool_two_slow": {"my_key": " slow"}}, - "assistant_id": "a", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "assistant_id": "a", + "thread_id": "1", + }, + parent_config=[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ + -1 + ].config, + interrupts=(), + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread1, debug=1) == { + "my_key": "value slow", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value slow", "market": "DE"}, + tasks=(), + next=(), + config={ + "configurable": { "thread_id": "1", - }, - parent_config=[ - c async for c in tool_two.checkpointer.alist(thread1, limit=2) - ][-1].config, - interrupts=(), - ) + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"tool_two_slow": {"my_key": " slow"}}, + "assistant_id": "a", + "thread_id": "1", + }, + parent_config=[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ + -1 + ].config, + interrupts=(), + ) - thread2 = {"configurable": {"thread_id": "2", "assistant_id": "a"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { - "my_key": "value", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), - next=("tool_two_fast",), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "assistant_id": "a", + thread2 = {"configurable": {"thread_id": "2", "assistant_id": "a"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), + next=("tool_two_fast",), + config={ + "configurable": { "thread_id": "2", - }, - parent_config=([c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ - -1 - ].config - ), - interrupts=(), - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread2, debug=1) == { - "my_key": "value fast", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value fast", "market": "US"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"tool_two_fast": {"my_key": " fast"}}, - "assistant_id": "a", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "assistant_id": "a", + "thread_id": "2", + }, + parent_config=( + [c async for c in tool_two.checkpointer.alist(thread2, limit=2)][-1].config + ), + interrupts=(), + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread2, debug=1) == { + "my_key": "value fast", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value fast", "market": "US"}, + tasks=(), + next=(), + config={ + "configurable": { "thread_id": "2", - }, - parent_config=([c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ - -1 - ].config - ), - interrupts=(), - ) + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"tool_two_fast": {"my_key": " fast"}}, + "assistant_id": "a", + "thread_id": "2", + }, + parent_config=( + [c async for c in tool_two.checkpointer.alist(thread2, limit=2)][-1].config + ), + interrupts=(), + ) - thread3 = {"configurable": {"thread_id": "3", "assistant_id": "b"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread3) == { - "my_key": "value", - "market": "US", - } - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "value", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), - next=("tool_two_fast",), - config={ - "configurable": { - "thread_id": "3", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "assistant_id": "b", + thread3 = {"configurable": {"thread_id": "3", "assistant_id": "b"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread3) == { + "my_key": "value", + "market": "US", + } + assert await tool_two.aget_state(thread3) == StateSnapshot( + values={"my_key": "value", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), + next=("tool_two_fast",), + config={ + "configurable": { "thread_id": "3", - }, - parent_config=([c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ - -1 - ].config - ), - interrupts=(), - ) - # update state - 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", (PULL, "tool_two_fast")),), - next=("tool_two_fast",), - config={ - "configurable": { - "thread_id": "3", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": {START: {"my_key": "key"}}, - "assistant_id": "b", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "assistant_id": "b", + "thread_id": "3", + }, + parent_config=( + [c async for c in tool_two.checkpointer.alist(thread3, limit=2)][-1].config + ), + interrupts=(), + ) + # update state + 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", (PULL, "tool_two_fast")),), + next=("tool_two_fast",), + config={ + "configurable": { "thread_id": "3", - }, - parent_config=([c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ - -1 - ].config - ), - interrupts=(), - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread3, debug=1) == { - "my_key": "valuekey fast", - "market": "US", - } - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "valuekey fast", "market": "US"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "3", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 2, - "writes": {"tool_two_fast": {"my_key": " fast"}}, - "assistant_id": "b", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": {START: {"my_key": "key"}}, + "assistant_id": "b", + "thread_id": "3", + }, + parent_config=( + [c async for c in tool_two.checkpointer.alist(thread3, limit=2)][-1].config + ), + interrupts=(), + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread3, debug=1) == { + "my_key": "valuekey fast", + "market": "US", + } + assert await tool_two.aget_state(thread3) == StateSnapshot( + values={"my_key": "valuekey fast", "market": "US"}, + tasks=(), + next=(), + config={ + "configurable": { "thread_id": "3", - }, - parent_config=([c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ - -1 - ].config - ), - interrupts=(), - ) + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 2, + "writes": {"tool_two_fast": {"my_key": " fast"}}, + "assistant_id": "b", + "thread_id": "3", + }, + parent_config=( + [c async for c in tool_two.checkpointer.alist(thread3, limit=2)][-1].config + ), + interrupts=(), + ) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_branch_then(checkpointer_name: str) -> None: +async def test_branch_then(async_checkpointer: BaseCheckpointSaver) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -4236,824 +4220,811 @@ async def test_branch_then(checkpointer_name: str) -> None: "market": "US", } - async with awith_checkpointer(checkpointer_name) as checkpointer: - # test stream_mode=debug - tool_two = tool_two_graph.compile(checkpointer=checkpointer) - thread10 = {"configurable": {"thread_id": "10"}} - assert [ - c - async for c in tool_two.astream( - {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" - ) - ] == [ - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": -1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": {"my_key": ""}, - "metadata": { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value", "market": "DE"}}, - "thread_id": "10", - }, - "parent_config": None, - "next": ["__start__"], - "tasks": [ - { - "id": AnyStr(), - "name": "__start__", - "interrupts": (), - "state": None, - } - ], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "10", - }, - "parent_config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "next": ["prepare"], - "tasks": [ - { - "id": AnyStr(), - "name": "prepare", - "interrupts": (), - "state": None, - } - ], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "prepare", - "input": {"my_key": "value", "market": "DE"}, - "triggers": ("branch:to:prepare",), - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "prepare", - "result": [("my_key", " prepared")], - "error": None, - "interrupts": [], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value prepared", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - "thread_id": "10", - }, - "parent_config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "next": ["tool_two_slow"], - "tasks": [ - { - "id": AnyStr(), - "name": "tool_two_slow", - "interrupts": (), - "state": None, - } - ], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": AnyStr(), - "name": "tool_two_slow", - "input": {"my_key": "value prepared", "market": "DE"}, - "triggers": ("branch:to:tool_two_slow",), - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "id": AnyStr(), - "name": "tool_two_slow", - "result": [("my_key", " slow")], - "error": None, - "interrupts": [], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 2, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value prepared slow", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 2, - "writes": {"tool_two_slow": {"my_key": " slow"}}, - "thread_id": "10", - }, - "parent_config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "next": ["finish"], - "tasks": [ - { - "id": AnyStr(), - "name": "finish", - "interrupts": (), - "state": None, - } - ], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": AnyStr(), - "name": "finish", - "input": {"my_key": "value prepared slow", "market": "DE"}, - "triggers": ( - "branch:prepare:condition::then", - "branch:to:finish", - ), - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "id": AnyStr(), - "name": "finish", - "result": [("my_key", " finished")], - "error": None, - "interrupts": [], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 3, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value prepared slow finished", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - "thread_id": "10", - }, - "parent_config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "10", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "next": [], - "tasks": [], - }, - }, - ] - - tool_two = tool_two_graph.compile( - checkpointer=checkpointer, - interrupt_before=["tool_two_fast", "tool_two_slow"], + # test stream_mode=debug + tool_two = tool_two_graph.compile(checkpointer=async_checkpointer) + thread10 = {"configurable": {"thread_id": "10"}} + assert [ + c + async for c in tool_two.astream( + {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" ) + ] == [ + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": -1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": {"my_key": ""}, + "metadata": { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value", "market": "DE"}}, + "thread_id": "10", + }, + "parent_config": None, + "next": ["__start__"], + "tasks": [ + { + "id": AnyStr(), + "name": "__start__", + "interrupts": (), + "state": None, + } + ], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value", + "market": "DE", + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "10", + }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "next": ["prepare"], + "tasks": [ + { + "id": AnyStr(), + "name": "prepare", + "interrupts": (), + "state": None, + } + ], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": AnyStr(), + "name": "prepare", + "input": {"my_key": "value", "market": "DE"}, + "triggers": ("branch:to:prepare",), + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": AnyStr(), + "name": "prepare", + "result": [("my_key", " prepared")], + "error": None, + "interrupts": [], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value prepared", + "market": "DE", + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "10", + }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "next": ["tool_two_slow"], + "tasks": [ + { + "id": AnyStr(), + "name": "tool_two_slow", + "interrupts": (), + "state": None, + } + ], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "tool_two_slow", + "input": {"my_key": "value prepared", "market": "DE"}, + "triggers": ("branch:to:tool_two_slow",), + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "id": AnyStr(), + "name": "tool_two_slow", + "result": [("my_key", " slow")], + "error": None, + "interrupts": [], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 2, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value prepared slow", + "market": "DE", + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 2, + "writes": {"tool_two_slow": {"my_key": " slow"}}, + "thread_id": "10", + }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "next": ["finish"], + "tasks": [ + { + "id": AnyStr(), + "name": "finish", + "interrupts": (), + "state": None, + } + ], + }, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": AnyStr(), + "name": "finish", + "input": {"my_key": "value prepared slow", "market": "DE"}, + "triggers": ( + "branch:prepare:condition::then", + "branch:to:finish", + ), + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "id": AnyStr(), + "name": "finish", + "result": [("my_key", " finished")], + "error": None, + "interrupts": [], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 3, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value prepared slow finished", + "market": "DE", + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "10", + }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "10"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "10", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "next": [], + "tasks": [], + }, + }, + ] - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + tool_two = tool_two_graph.compile( + checkpointer=async_checkpointer, + interrupt_before=["tool_two_fast", "tool_two_slow"], + ) - thread1 = {"configurable": {"thread_id": "11"}} - # stop when about to enter node - assert [ - c - async for c in tool_two.astream( - {"my_key": "value", "market": "DE"}, thread1, stream_mode="debug" - ) - ] == [ - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": -1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": {"my_key": ""}, - "metadata": { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value", "market": "DE"}}, + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + + thread1 = {"configurable": {"thread_id": "11"}} + # stop when about to enter node + assert [ + c + async for c in tool_two.astream( + {"my_key": "value", "market": "DE"}, thread1, stream_mode="debug" + ) + ] == [ + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": -1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "11"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { "thread_id": "11", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), }, - "parent_config": None, - "next": ["__start__"], - "tasks": [ - { - "id": AnyStr(), - "name": "__start__", - "interrupts": (), - "state": None, - } - ], }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "11", - }, - "parent_config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "next": ["prepare"], - "tasks": [ - { - "id": AnyStr(), - "name": "prepare", - "interrupts": (), - "state": None, - } - ], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "prepare", - "input": {"my_key": "value", "market": "DE"}, - "triggers": ("branch:to:prepare",), - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "prepare", - "result": [("my_key", " prepared")], - "error": None, - "interrupts": [], - }, - }, - { - "type": "checkpoint", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "values": { - "my_key": "value prepared", - "market": "DE", - }, - "metadata": { - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - "thread_id": "11", - }, - "parent_config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, - "configurable": { - "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - }, - }, - "next": ["tool_two_slow"], - "tasks": [ - { - "id": AnyStr(), - "name": "tool_two_slow", - "interrupts": (), - "state": None, - } - ], - }, - }, - ] - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), - next=("tool_two_slow",), - config={ - "configurable": { + "values": {"my_key": ""}, + "metadata": { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value", "market": "DE"}}, "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } + }, + "parent_config": None, + "next": ["__start__"], + "tasks": [ + { + "id": AnyStr(), + "name": "__start__", + "interrupts": (), + "state": None, + } + ], }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, - "thread_id": "11", - }, - parent_config=([c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config - ), - interrupts=(), - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread1, debug=1) == { - "my_key": "value prepared slow finished", - "market": "DE", - } - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value prepared slow finished", "market": "DE"}, - tasks=(), - next=(), - config={ - "configurable": { + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 0, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "11"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "11", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value", + "market": "DE", + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, "thread_id": "11", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } + }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "11"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "11", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "next": ["prepare"], + "tasks": [ + { + "id": AnyStr(), + "name": "prepare", + "interrupts": (), + "state": None, + } + ], }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, + }, + { + "type": "task", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": AnyStr(), + "name": "prepare", + "input": {"my_key": "value", "market": "DE"}, + "triggers": ("branch:to:prepare",), + }, + }, + { + "type": "task_result", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "id": AnyStr(), + "name": "prepare", + "result": [("my_key", " prepared")], + "error": None, + "interrupts": [], + }, + }, + { + "type": "checkpoint", + "timestamp": AnyStr(), + "step": 1, + "payload": { + "config": { + "tags": [], + "metadata": {"thread_id": "11"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "11", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "values": { + "my_key": "value prepared", + "market": "DE", + }, + "metadata": { + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "11", + }, + "parent_config": { + "tags": [], + "metadata": {"thread_id": "11"}, + "callbacks": None, + "recursion_limit": 25, + "configurable": { + "thread_id": "11", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + }, + }, + "next": ["tool_two_slow"], + "tasks": [ + { + "id": AnyStr(), + "name": "tool_two_slow", + "interrupts": (), + "state": None, + } + ], + }, + }, + ] + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), + next=("tool_two_slow",), + config={ + "configurable": { "thread_id": "11", - }, - parent_config=([c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config - ), - interrupts=(), - ) + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "11", + }, + parent_config=( + [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][-1].config + ), + interrupts=(), + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread1, debug=1) == { + "my_key": "value prepared slow finished", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value prepared slow finished", "market": "DE"}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "11", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "11", + }, + parent_config=( + [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][-1].config + ), + interrupts=(), + ) - thread2 = {"configurable": {"thread_id": "12"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { - "my_key": "value prepared", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value prepared", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), - next=("tool_two_fast",), - config={ - "configurable": { - "thread_id": "12", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, + thread2 = {"configurable": {"thread_id": "12"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value prepared", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value prepared", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), + next=("tool_two_fast",), + config={ + "configurable": { "thread_id": "12", - }, - parent_config=([c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ - -1 - ].config - ), - interrupts=(), - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread2, debug=1) == { - "my_key": "value prepared fast finished", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value prepared fast finished", "market": "US"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "12", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "12", + }, + parent_config=( + [c async for c in tool_two.checkpointer.alist(thread2, limit=2)][-1].config + ), + interrupts=(), + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread2, debug=1) == { + "my_key": "value prepared fast finished", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value prepared fast finished", "market": "US"}, + tasks=(), + next=(), + config={ + "configurable": { "thread_id": "12", - }, - parent_config=([c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ - -1 - ].config - ), - interrupts=(), - ) + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "12", + }, + parent_config=( + [c async for c in tool_two.checkpointer.alist(thread2, limit=2)][-1].config + ), + interrupts=(), + ) - tool_two = tool_two_graph.compile( - checkpointer=checkpointer, interrupt_after=["prepare"] - ) + tool_two = tool_two_graph.compile( + checkpointer=async_checkpointer, interrupt_after=["prepare"] + ) - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - thread1 = {"configurable": {"thread_id": "21"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == { - "my_key": "value prepared", - "market": "DE", - } - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), - next=("tool_two_slow",), - config={ - "configurable": { - "thread_id": "21", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, + thread1 = {"configurable": {"thread_id": "21"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "DE"}, thread1) == { + "my_key": "value prepared", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), + next=("tool_two_slow",), + config={ + "configurable": { "thread_id": "21", - }, - parent_config=([c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config - ), - interrupts=(), - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread1, debug=1) == { - "my_key": "value prepared slow finished", - "market": "DE", - } - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value prepared slow finished", "market": "DE"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "21", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "21", + }, + parent_config=( + [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][-1].config + ), + interrupts=(), + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread1, debug=1) == { + "my_key": "value prepared slow finished", + "market": "DE", + } + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value prepared slow finished", "market": "DE"}, + tasks=(), + next=(), + config={ + "configurable": { "thread_id": "21", - }, - parent_config=([c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config - ), - interrupts=(), - ) + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "21", + }, + parent_config=( + [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][-1].config + ), + interrupts=(), + ) - thread2 = {"configurable": {"thread_id": "22"}} - # stop when about to enter node - assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { - "my_key": "value prepared", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value prepared", "market": "US"}, - tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), - next=("tool_two_fast",), - config={ - "configurable": { - "thread_id": "22", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, + thread2 = {"configurable": {"thread_id": "22"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}, thread2) == { + "my_key": "value prepared", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value prepared", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast", (PULL, "tool_two_fast")),), + next=("tool_two_fast",), + config={ + "configurable": { "thread_id": "22", - }, - parent_config=([c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ - -1 - ].config - ), - interrupts=(), - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread2, debug=1) == { - "my_key": "value prepared fast finished", - "market": "US", - } - assert await tool_two.aget_state(thread2) == StateSnapshot( - values={"my_key": "value prepared fast finished", "market": "US"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "22", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "22", + }, + parent_config=( + [c async for c in tool_two.checkpointer.alist(thread2, limit=2)][-1].config + ), + interrupts=(), + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread2, debug=1) == { + "my_key": "value prepared fast finished", + "market": "US", + } + assert await tool_two.aget_state(thread2) == StateSnapshot( + values={"my_key": "value prepared fast finished", "market": "US"}, + tasks=(), + next=(), + config={ + "configurable": { "thread_id": "22", - }, - parent_config=([c async for c in tool_two.checkpointer.alist(thread2, limit=2)][ - -1 - ].config - ), - interrupts=(), - ) + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "22", + }, + parent_config=( + [c async for c in tool_two.checkpointer.alist(thread2, limit=2)][-1].config + ), + interrupts=(), + ) - thread3 = {"configurable": {"thread_id": "23"}} - # update an empty thread before first run - uconfig = await tool_two.aupdate_state( - thread3, {"my_key": "key", "market": "DE"} - ) - # check current state - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "key", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "prepare", (PULL, "prepare")),), - next=("prepare",), - config=uconfig, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "update", - "step": 0, - "writes": {START: {"my_key": "key", "market": "DE"}}, + thread3 = {"configurable": {"thread_id": "23"}} + # update an empty thread before first run + uconfig = await tool_two.aupdate_state(thread3, {"my_key": "key", "market": "DE"}) + # check current state + assert await tool_two.aget_state(thread3) == StateSnapshot( + values={"my_key": "key", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "prepare", (PULL, "prepare")),), + next=("prepare",), + config=uconfig, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "update", + "step": 0, + "writes": {START: {"my_key": "key", "market": "DE"}}, + "thread_id": "23", + }, + parent_config=None, + interrupts=(), + ) + # run from this point + assert await tool_two.ainvoke(None, thread3) == { + "my_key": "key prepared", + "market": "DE", + } + # get state after first node + assert await tool_two.aget_state(thread3) == StateSnapshot( + values={"my_key": "key prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), + next=("tool_two_slow",), + config={ + "configurable": { "thread_id": "23", - }, - parent_config=None, - interrupts=(), - ) - # run from this point - assert await tool_two.ainvoke(None, thread3) == { - "my_key": "key prepared", - "market": "DE", - } - # get state after first node - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "key prepared", "market": "DE"}, - tasks=(PregelTask(AnyStr(), "tool_two_slow", (PULL, "tool_two_slow")),), - next=("tool_two_slow",), - config={ - "configurable": { - "thread_id": "23", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "writes": {"prepare": {"my_key": " prepared"}}, + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "23", + }, + parent_config=(uconfig), + interrupts=(), + ) + # resume, for same result as above + assert await tool_two.ainvoke(None, thread3, debug=1) == { + "my_key": "key prepared slow finished", + "market": "DE", + } + assert await tool_two.aget_state(thread3) == StateSnapshot( + values={"my_key": "key prepared slow finished", "market": "DE"}, + tasks=(), + next=(), + config={ + "configurable": { "thread_id": "23", - }, - parent_config=(uconfig), - interrupts=(), - ) - # resume, for same result as above - assert await tool_two.ainvoke(None, thread3, debug=1) == { - "my_key": "key prepared slow finished", - "market": "DE", - } - assert await tool_two.aget_state(thread3) == StateSnapshot( - values={"my_key": "key prepared slow finished", "market": "DE"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "23", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 3, - "writes": {"finish": {"my_key": " finished"}}, - "thread_id": "23", - }, - parent_config=([c async for c in tool_two.checkpointer.alist(thread3, limit=2)][ - -1 - ].config - ), - interrupts=(), - ) + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "23", + }, + parent_config=( + [c async for c in tool_two.checkpointer.alist(thread3, limit=2)][-1].config + ), + interrupts=(), + ) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_nested_graph_state(checkpointer_name: str) -> None: +async def test_nested_graph_state(async_checkpointer: BaseCheckpointSaver) -> None: class InnerState(TypedDict): my_key: str my_other_key: str @@ -5099,291 +5070,71 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: graph.add_edge("inner", "outer_2") graph.set_finish_point("outer_2") - async with awith_checkpointer(checkpointer_name) as checkpointer: - app = graph.compile(checkpointer=checkpointer) + app = graph.compile(checkpointer=async_checkpointer) - config = {"configurable": {"thread_id": "1"}} - await app.ainvoke({"my_key": "my value"}, config, debug=True) - # test state w/ nested subgraph state (right after interrupt) - # first get_state without subgraph state - assert await app.aget_state(config) == StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - (PULL, "inner"), - state={ - "configurable": {"thread_id": "1", "checkpoint_ns": AnyStr()} - }, - ), + config = {"configurable": {"thread_id": "1"}} + await app.ainvoke({"my_key": "my value"}, config, debug=True) + # test state w/ nested subgraph state (right after interrupt) + # first get_state without subgraph state + assert await app.aget_state(config) == StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "inner", + (PULL, "inner"), + state={"configurable": {"thread_id": "1", "checkpoint_ns": AnyStr()}}, ), - next=("inner",), - config={ + ), + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": "", "checkpoint_id": AnyStr(), } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=({ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - interrupts=(), - ) - # now, get_state with subgraphs state - assert await app.aget_state(config, subgraphs=True) == StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - (PULL, "inner"), - state=StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=( - PregelTask( - AnyStr(), - "inner_2", - (PULL, "inner_2"), - ), - ), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "parents": { - "": AnyStr(), - }, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - "thread_id": "1", - "langgraph_node": "inner", - "langgraph_path": [PULL, "inner"], - "langgraph_step": 2, - "langgraph_triggers": ["branch:to:inner"], - "langgraph_checkpoint_ns": AnyStr("inner:"), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - interrupts=(), - ), - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=({ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - interrupts=(), - ) - # get_state_history returns outer graph checkpoints - history = [c async for c in app.aget_state_history(config)] - expected_history = [ - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - (PULL, "inner"), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - } - }, - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=({ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - interrupts=(), - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=( - PregelTask( - AnyStr(), - "outer_1", - (PULL, "outer_1"), - result={"my_key": "hi my value"}, - ), - ), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": None, - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - interrupts=(), - ), - StateSnapshot( - values={}, - tasks=( - PregelTask( - AnyStr(), - "__start__", - (PULL, "__start__"), - result={"my_key": "my value"}, - ), - ), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"__start__": {"my_key": "my value"}}, - "step": -1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - interrupts=(), - ), - ] - - assert history == expected_history - - # get_state_history for a subgraph returns its checkpoints - child_history = [ - c async for c in app.aget_state_history(history[0].tasks[0].state) - ] - expected_child_history = [ - StateSnapshot( - values={"my_key": "hi my value here", "my_other_key": "hi my value"}, - next=("inner_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } + } + ), + interrupts=(), + ) + # now, get_state with subgraphs state + assert await app.aget_state(config, subgraphs=True) == StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "inner", + (PULL, "inner"), + state=StateSnapshot( + values={ + "my_key": "hi my value here", + "my_other_key": "hi my value", }, - "step": 1, - "parents": {"": AnyStr()}, - "thread_id": "1", - "langgraph_node": "inner", - "langgraph_path": [PULL, "inner"], - "langgraph_step": 2, - "langgraph_triggers": ["branch:to:inner"], - "langgraph_checkpoint_ns": AnyStr("inner:"), - }, - created_at=AnyStr(), - parent_config=({ + tasks=( + PregelTask( + AnyStr(), + "inner_2", + (PULL, "inner_2"), + ), + ), + next=("inner_2",), + config={ "configurable": { "thread_id": "1", "checkpoint_ns": AnyStr("inner:"), @@ -5392,64 +5143,217 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: {"": AnyStr(), AnyStr("child:"): AnyStr()} ), } - } - ), - tasks=(PregelTask(AnyStr(), "inner_2", (PULL, "inner_2")),), - interrupts=(), - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "loop", - "writes": None, - "step": 0, - "parents": {"": AnyStr()}, - "thread_id": "1", - "langgraph_node": "inner", - "langgraph_path": [PULL, "inner"], - "langgraph_step": 2, - "langgraph_triggers": ["branch:to:inner"], - "langgraph_checkpoint_ns": AnyStr("inner:"), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - tasks=( - PregelTask( - AnyStr(), - "inner_1", - (PULL, "inner_1"), - result={ - "my_key": "hi my value here", - "my_other_key": "hi my value", + }, + metadata={ + "parents": { + "": AnyStr(), }, - ), + "source": "loop", + "writes": { + "inner_1": { + "my_key": "hi my value here", + "my_other_key": "hi my value", + } + }, + "step": 1, + "thread_id": "1", + "langgraph_node": "inner", + "langgraph_path": [PULL, "inner"], + "langgraph_step": 2, + "langgraph_triggers": ["branch:to:inner"], + "langgraph_checkpoint_ns": AnyStr("inner:"), + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + interrupts=(), ), - interrupts=(), ), - StateSnapshot( - values={}, - next=("__start__",), - config={ + ), + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=( + { + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + } + ), + interrupts=(), + ) + # get_state_history returns outer graph checkpoints + history = [c async for c in app.aget_state_history(config)] + expected_history = [ + StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "inner", + (PULL, "inner"), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + } + }, + ), + ), + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=( + { + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + } + ), + interrupts=(), + ), + StateSnapshot( + values={"my_key": "my value"}, + tasks=( + PregelTask( + AnyStr(), + "outer_1", + (PULL, "outer_1"), + result={"my_key": "hi my value"}, + ), + ), + next=("outer_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": None, + "step": 0, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + interrupts=(), + ), + StateSnapshot( + values={}, + tasks=( + PregelTask( + AnyStr(), + "__start__", + (PULL, "__start__"), + result={"my_key": "my value"}, + ), + ), + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "writes": {"__start__": {"my_key": "my value"}}, + "step": -1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ), + ] + + assert history == expected_history + + # get_state_history for a subgraph returns its checkpoints + child_history = [c async for c in app.aget_state_history(history[0].tasks[0].state)] + expected_child_history = [ + StateSnapshot( + values={"my_key": "hi my value here", "my_other_key": "hi my value"}, + next=("inner_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "source": "loop", + "writes": { + "inner_1": { + "my_key": "hi my value here", + "my_other_key": "hi my value", + } + }, + "step": 1, + "parents": {"": AnyStr()}, + "thread_id": "1", + "langgraph_node": "inner", + "langgraph_path": [PULL, "inner"], + "langgraph_step": 2, + "langgraph_triggers": ["branch:to:inner"], + "langgraph_checkpoint_ns": AnyStr("inner:"), + }, + created_at=AnyStr(), + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": AnyStr("inner:"), @@ -5458,39 +5362,140 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: {"": AnyStr(), AnyStr("child:"): AnyStr()} ), } - }, - metadata={ - "source": "input", - "writes": {"__start__": {"my_key": "hi my value"}}, - "step": -1, - "parents": {"": AnyStr()}, - "thread_id": "1", - "langgraph_node": "inner", - "langgraph_path": [PULL, "inner"], - "langgraph_step": 2, - "langgraph_triggers": ["branch:to:inner"], - "langgraph_checkpoint_ns": AnyStr("inner:"), - }, - created_at=AnyStr(), - parent_config=None, - tasks=( - PregelTask( - AnyStr(), - "__start__", - (PULL, "__start__"), - result={"my_key": "hi my value"}, - ), - ), - interrupts=(), + } ), - ] + tasks=(PregelTask(AnyStr(), "inner_2", (PULL, "inner_2")),), + interrupts=(), + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "source": "loop", + "writes": None, + "step": 0, + "parents": {"": AnyStr()}, + "thread_id": "1", + "langgraph_node": "inner", + "langgraph_path": [PULL, "inner"], + "langgraph_step": 2, + "langgraph_triggers": ["branch:to:inner"], + "langgraph_checkpoint_ns": AnyStr("inner:"), + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + tasks=( + PregelTask( + AnyStr(), + "inner_1", + (PULL, "inner_1"), + result={ + "my_key": "hi my value here", + "my_other_key": "hi my value", + }, + ), + ), + interrupts=(), + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "source": "input", + "writes": {"__start__": {"my_key": "hi my value"}}, + "step": -1, + "parents": {"": AnyStr()}, + "thread_id": "1", + "langgraph_node": "inner", + "langgraph_path": [PULL, "inner"], + "langgraph_step": 2, + "langgraph_triggers": ["branch:to:inner"], + "langgraph_checkpoint_ns": AnyStr("inner:"), + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + AnyStr(), + "__start__", + (PULL, "__start__"), + result={"my_key": "hi my value"}, + ), + ), + interrupts=(), + ), + ] - assert child_history == expected_child_history + assert child_history == expected_child_history - # resume - await app.ainvoke(None, config, debug=True) - # test state w/ nested subgraph state (after resuming from interrupt) - assert await app.aget_state(config) == StateSnapshot( + # resume + await app.ainvoke(None, config, debug=True) + # test state w/ nested subgraph state (after resuming from interrupt) + assert await app.aget_state(config) == StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": { + "outer_2": {"my_key": "hi my value here and there and back again"} + }, + "step": 3, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=( + { + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + } + ), + interrupts=(), + ) + # test full history at the end + actual_history = [c async for c in app.aget_state_history(config)] + expected_history = [ + StateSnapshot( values={"my_key": "hi my value here and there and back again"}, tasks=(), next=(), @@ -5511,7 +5516,8 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -5520,193 +5526,158 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: } ), interrupts=(), - ) - # test full history at the end - actual_history = [c async for c in app.aget_state_history(config)] - expected_history = [ - StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "outer_2": { - "my_key": "hi my value here and there and back again" - } - }, - "step": 3, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + tasks=( + PregelTask( + AnyStr(), + "outer_2", + (PULL, "outer_2"), + result={"my_key": "hi my value here and there and back again"}, + ), + ), + next=("outer_2",), + config={ + "configurable": { "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=({ + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"inner": {"my_key": "hi my value here and there"}}, + "step": 2, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + interrupts=(), + ), + StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "inner", + (PULL, "inner"), + state={ "configurable": { "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), + "checkpoint_ns": AnyStr(), } - } + }, + result={"my_key": "hi my value here and there"}, ), - interrupts=(), ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - tasks=( - PregelTask( - AnyStr(), - "outer_2", - (PULL, "outer_2"), - result={"my_key": "hi my value here and there and back again"}, - ), - ), - next=("outer_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"inner": {"my_key": "hi my value here and there"}}, - "step": 2, + next=("inner",), + config={ + "configurable": { "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - interrupts=(), - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - (PULL, "inner"), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - } - }, - result={"my_key": "hi my value here and there"}, - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - interrupts=(), - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=( - PregelTask( - AnyStr(), - "outer_1", - (PULL, "outer_1"), - result={"my_key": "hi my value"}, - ), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + interrupts=(), + ), + StateSnapshot( + values={"my_key": "my value"}, + tasks=( + PregelTask( + AnyStr(), + "outer_1", + (PULL, "outer_1"), + result={"my_key": "hi my value"}, ), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": None, - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - interrupts=(), ), - StateSnapshot( - values={}, - tasks=( - PregelTask( - AnyStr(), - "__start__", - (PULL, "__start__"), - result={"my_key": "my value"}, - ), + next=("outer_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": None, + "step": 0, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + interrupts=(), + ), + StateSnapshot( + values={}, + tasks=( + PregelTask( + AnyStr(), + "__start__", + (PULL, "__start__"), + result={"my_key": "my value"}, ), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"__start__": {"my_key": "my value"}}, - "step": -1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - interrupts=(), ), - ] + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "writes": {"__start__": {"my_key": "my value"}}, + "step": -1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ), + ] - assert actual_history == expected_history - # test looking up parent state by checkpoint ID - for actual_snapshot, expected_snapshot in zip(actual_history, expected_history): - assert await app.aget_state(actual_snapshot.config) == expected_snapshot + assert actual_history == expected_history + # test looking up parent state by checkpoint ID + for actual_snapshot, expected_snapshot in zip(actual_history, expected_history): + assert await app.aget_state(actual_snapshot.config) == expected_snapshot -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: +async def test_doubly_nested_graph_state( + async_checkpointer: BaseCheckpointSaver, +) -> None: class State(TypedDict): my_key: str @@ -5754,82 +5725,109 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: graph.add_edge("child", "parent_2") graph.set_finish_point("parent_2") - async with awith_checkpointer(checkpointer_name) as checkpointer: - app = graph.compile(checkpointer=checkpointer) + app = graph.compile(checkpointer=async_checkpointer) - # test invoke w/ nested interrupt - config = {"configurable": {"thread_id": "1"}} - assert [ - c async for c in app.astream({"my_key": "my value"}, config, subgraphs=True) - ] == [ - ((), {"parent_1": {"my_key": "hi my value"}}), - ( - (AnyStr("child:"), AnyStr("child_1:")), - {"grandchild_1": {"my_key": "hi my value here"}}, + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert [ + c async for c in app.astream({"my_key": "my value"}, config, subgraphs=True) + ] == [ + ((), {"parent_1": {"my_key": "hi my value"}}), + ( + (AnyStr("child:"), AnyStr("child_1:")), + {"grandchild_1": {"my_key": "hi my value here"}}, + ), + ((), {"__interrupt__": ()}), + ] + # get state without subgraphs + outer_state = await app.aget_state(config) + assert outer_state == StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child", + (PULL, "child"), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child"), + } + }, ), - ((), {"__interrupt__": ()}), - ] - # get state without subgraphs - outer_state = await app.aget_state(config) - assert outer_state == StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "child", - (PULL, "child"), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child"), - } - }, - ), - ), - next=("child",), - config={ + ), + next=("child",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"parent_1": {"my_key": "hi my value"}}, + "step": 1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": "", "checkpoint_id": AnyStr(), } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"parent_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=({ + } + ), + interrupts=(), + ) + child_state = await app.aget_state(outer_state.tasks[0].state) + assert child_state == StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child_1", + (PULL, "child_1"), + state={ "configurable": { "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), + "checkpoint_ns": AnyStr(), } - } + }, ), - interrupts=(), - ) - child_state = await app.aget_state(outer_state.tasks[0].state) - assert child_state == StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "child_1", - (PULL, "child_1"), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - } - }, + ), + next=("child_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } ), - ), - next=("child_1",), - config={ + } + }, + metadata={ + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child", + "langgraph_path": ["__pregel_pull", "child"], + "langgraph_step": 2, + "langgraph_triggers": ["branch:to:child"], + "parents": {"": AnyStr()}, + "source": "loop", + "writes": None, + "step": 0, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": AnyStr("child:"), @@ -5841,48 +5839,57 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: } ), } - }, - metadata={ - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child", - "langgraph_path": ["__pregel_pull", "child"], - "langgraph_step": 2, - "langgraph_triggers": ["branch:to:child"], - "parents": {"": AnyStr()}, - "source": "loop", - "writes": None, - "step": 0, + } + ), + interrupts=(), + ) + grandchild_state = await app.aget_state(child_state.tasks[0].state) + assert grandchild_state == StateSnapshot( + values={"my_key": "hi my value here"}, + tasks=( + PregelTask( + AnyStr(), + "grandchild_2", + (PULL, "grandchild_2"), + ), + ), + next=("grandchild_2",), + config={ + "configurable": { "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=({ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), } + ), + } + }, + metadata={ + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), } ), - interrupts=(), - ) - grandchild_state = await app.aget_state(child_state.tasks[0].state) - assert grandchild_state == StateSnapshot( - values={"my_key": "hi my value here"}, - tasks=( - PregelTask( - AnyStr(), - "grandchild_2", - (PULL, "grandchild_2"), - ), - ), - next=("grandchild_2",), - config={ + "source": "loop", + "writes": {"grandchild_1": {"my_key": "hi my value here"}}, + "step": 1, + "thread_id": "1", + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [PULL, AnyStr("child_1")], + "langgraph_step": 1, + "langgraph_triggers": [ + "branch:to:child_1", + ], + }, + created_at=AnyStr(), + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": AnyStr(), @@ -5895,70 +5902,78 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: } ), } - }, - metadata={ - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "source": "loop", - "writes": {"grandchild_1": {"my_key": "hi my value here"}}, - "step": 1, - "thread_id": "1", - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [PULL, AnyStr("child_1")], - "langgraph_step": 1, - "langgraph_triggers": [ - "branch:to:child_1", - ], - }, - created_at=AnyStr(), - parent_config=({ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - } - ), - interrupts=(), - ) - # get state with subgraphs - assert await app.aget_state(config, subgraphs=True) == StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "child", - (PULL, "child"), - state=StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "child_1", - (PULL, "child_1"), - state=StateSnapshot( - values={"my_key": "hi my value here"}, - tasks=( - PregelTask( - AnyStr(), - "grandchild_2", - (PULL, "grandchild_2"), - ), + } + ), + interrupts=(), + ) + # get state with subgraphs + assert await app.aget_state(config, subgraphs=True) == StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child", + (PULL, "child"), + state=StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child_1", + (PULL, "child_1"), + state=StateSnapshot( + values={"my_key": "hi my value here"}, + tasks=( + PregelTask( + AnyStr(), + "grandchild_2", + (PULL, "grandchild_2"), ), - next=("grandchild_2",), - config={ + ), + next=("grandchild_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr( + re.compile(r"child:.+|child1:") + ): AnyStr(), + } + ), + } + }, + metadata={ + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + "source": "loop", + "writes": { + "grandchild_1": {"my_key": "hi my value here"} + }, + "step": 1, + "thread_id": "1", + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [ + "branch:to:child_1", + ], + }, + created_at=AnyStr(), + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": AnyStr(), @@ -5973,59 +5988,40 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: } ), } - }, - metadata={ - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "source": "loop", - "writes": { - "grandchild_1": { - "my_key": "hi my value here" - } - }, - "step": 1, - "thread_id": "1", - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": [ - "branch:to:child_1", - ], - }, - created_at=AnyStr(), - parent_config=({ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr( - re.compile( - r"child:.+|child1:" - ) - ): AnyStr(), - } - ), - } - } - ), - interrupts=(), + } ), + interrupts=(), ), ), - next=("child_1",), - config={ + ), + next=("child_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "parents": {"": AnyStr()}, + "source": "loop", + "writes": None, + "step": 0, + "thread_id": "1", + "langgraph_node": "child", + "langgraph_path": [PULL, AnyStr("child")], + "langgraph_step": 2, + "langgraph_triggers": [ + "branch:to:child", + ], + "langgraph_checkpoint_ns": AnyStr("child:"), + }, + created_at=AnyStr(), + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": AnyStr("child:"), @@ -6034,38 +6030,60 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: {"": AnyStr(), AnyStr("child:"): AnyStr()} ), } - }, - metadata={ - "parents": {"": AnyStr()}, - "source": "loop", - "writes": None, - "step": 0, - "thread_id": "1", - "langgraph_node": "child", - "langgraph_path": [PULL, AnyStr("child")], - "langgraph_step": 2, - "langgraph_triggers": [ - "branch:to:child", - ], - "langgraph_checkpoint_ns": AnyStr("child:"), - }, - created_at=AnyStr(), - parent_config=({ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - } - ), - interrupts=(), + } ), + interrupts=(), ), ), - next=("child",), + ), + next=("child",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"parent_1": {"my_key": "hi my value"}}, + "step": 1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=( + { + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + } + ), + interrupts=(), + ) + # resume + assert [c async for c in app.astream(None, config, subgraphs=True)] == [ + ( + (AnyStr("child:"), AnyStr("child_1:")), + {"grandchild_2": {"my_key": "hi my value here and there"}}, + ), + ( + (AnyStr("child:"),), + {"child_1": {"my_key": "hi my value here and there"}}, + ), + ((), {"child": {"my_key": "hi my value here and there"}}), + ((), {"parent_2": {"my_key": "hi my value here and there and back again"}}), + ] + # get state with and without subgraphs + assert ( + await app.aget_state(config) + == await app.aget_state(config, subgraphs=True) + == StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + tasks=(), + next=(), config={ "configurable": { "thread_id": "1", @@ -6076,12 +6094,15 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: metadata={ "parents": {}, "source": "loop", - "writes": {"parent_1": {"my_key": "hi my value"}}, - "step": 1, + "writes": { + "parent_2": {"my_key": "hi my value here and there and back again"} + }, + "step": 3, "thread_id": "1", }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -6091,24 +6112,14 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: ), interrupts=(), ) - # resume - assert [c async for c in app.astream(None, config, subgraphs=True)] == [ - ( - (AnyStr("child:"), AnyStr("child_1:")), - {"grandchild_2": {"my_key": "hi my value here and there"}}, - ), - ( - (AnyStr("child:"),), - {"child_1": {"my_key": "hi my value here and there"}}, - ), - ((), {"child": {"my_key": "hi my value here and there"}}), - ((), {"parent_2": {"my_key": "hi my value here and there and back again"}}), - ] - # get state with and without subgraphs - assert ( - await app.aget_state(config) - == await app.aget_state(config, subgraphs=True) - == StateSnapshot( + ) + + # get outer graph history + outer_history = [c async for c in app.aget_state_history(config)] + assert ( + outer_history[0] + == [ + StateSnapshot( values={"my_key": "hi my value here and there and back again"}, tasks=(), next=(), @@ -6131,511 +6142,112 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", }, created_at=AnyStr(), - parent_config=({ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } - ), + }, interrupts=(), - ) - ) - - - # get outer graph history - outer_history = [c async for c in app.aget_state_history(config)] - assert ( - outer_history[0] - == [ - StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "parent_2": { - "my_key": "hi my value here and there and back again" - } - }, - "step": 3, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - interrupts=(), - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=("parent_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"child": {"my_key": "hi my value here and there"}}, - "step": 2, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), name="parent_2", path=(PULL, "parent_2") - ), - ), - interrupts=(), - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "child", - (PULL, "child"), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child"), - } - }, - ), - ), - next=("child",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"parent_1": {"my_key": "hi my value"}}, - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - interrupts=(), - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("parent_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": None, - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), name="parent_1", path=(PULL, "parent_1") - ), - ), - interrupts=(), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - tasks=( - PregelTask( - id=AnyStr(), name="__start__", path=(PULL, "__start__") - ), - ), - interrupts=(), - ), - ][0] - ) - # get child graph history - child_history = [ - c async for c in app.aget_state_history(outer_history[2].tasks[0].state) - ] - assert child_history == [ + ), StateSnapshot( values={"my_key": "hi my value here and there"}, - next=(), + next=("parent_2",), config={ "configurable": { "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), + "checkpoint_ns": "", "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), } }, metadata={ + "parents": {}, "source": "loop", - "writes": {"child_1": {"my_key": "hi my value here and there"}}, - "step": 1, - "parents": {"": AnyStr()}, + "writes": {"child": {"my_key": "hi my value here and there"}}, + "step": 2, "thread_id": "1", - "langgraph_node": "child", - "langgraph_path": [PULL, AnyStr("child")], - "langgraph_step": 2, - "langgraph_triggers": ["branch:to:child"], - "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), parent_config={ "configurable": { "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), + "checkpoint_ns": "", "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), } }, - tasks=(), + tasks=( + PregelTask(id=AnyStr(), name="parent_2", path=(PULL, "parent_2")), + ), interrupts=(), ), StateSnapshot( values={"my_key": "hi my value"}, - next=("child_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "loop", - "writes": None, - "step": 0, - "parents": {"": AnyStr()}, - "thread_id": "1", - "langgraph_node": "child", - "langgraph_path": [PULL, AnyStr("child")], - "langgraph_step": 2, - "langgraph_triggers": ["branch:to:child"], - "langgraph_checkpoint_ns": AnyStr("child:"), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, tasks=( PregelTask( - id=AnyStr(), - name="child_1", - path=(PULL, "child_1"), + AnyStr(), + "child", + (PULL, "child"), state={ "configurable": { "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), + "checkpoint_ns": AnyStr("child"), } }, - result={"my_key": "hi my value here and there"}, ), ), - interrupts=(), - ), - StateSnapshot( - values={}, - next=("__start__",), + next=("child",), config={ "configurable": { "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), + "checkpoint_ns": "", "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "input", - "writes": {"__start__": {"my_key": "hi my value"}}, - "step": -1, - "parents": {"": AnyStr()}, - "thread_id": "1", - "langgraph_node": "child", - "langgraph_path": [PULL, AnyStr("child")], - "langgraph_step": 2, - "langgraph_triggers": ["branch:to:child"], - "langgraph_checkpoint_ns": AnyStr("child:"), - }, - created_at=AnyStr(), - parent_config=None, - tasks=( - PregelTask( - id=AnyStr(), - name="__start__", - path=(PULL, "__start__"), - result={"my_key": "hi my value"}, - ), - ), - interrupts=(), - ), - ] - # get grandchild graph history - grandchild_history = [ - c async for c in app.aget_state_history(child_history[1].tasks[0].state) - ] - assert grandchild_history == [ - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), } }, metadata={ + "parents": {}, "source": "loop", - "writes": { - "grandchild_2": {"my_key": "hi my value here and there"} - }, - "step": 2, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "thread_id": "1", - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": [ - "branch:to:child_1", - ], - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, - tasks=(), - interrupts=(), - ), - StateSnapshot( - values={"my_key": "hi my value here"}, - next=("grandchild_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, - metadata={ - "source": "loop", - "writes": {"grandchild_1": {"my_key": "hi my value here"}}, + "writes": {"parent_1": {"my_key": "hi my value"}}, "step": 1, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), "thread_id": "1", - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": [ - "branch:to:child_1", - ], }, created_at=AnyStr(), parent_config={ "configurable": { "thread_id": "1", - "checkpoint_ns": AnyStr(), + "checkpoint_ns": "", "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), } }, - tasks=( - PregelTask( - id=AnyStr(), - name="grandchild_2", - path=(PULL, "grandchild_2"), - result={"my_key": "hi my value here and there"}, - ), - ), interrupts=(), ), StateSnapshot( - values={"my_key": "hi my value"}, - next=("grandchild_1",), + values={"my_key": "my value"}, + next=("parent_1",), config={ "configurable": { "thread_id": "1", - "checkpoint_ns": AnyStr(), + "checkpoint_ns": "", "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), } }, metadata={ + "parents": {}, "source": "loop", "writes": None, "step": 0, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), "thread_id": "1", - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": [ - "branch:to:child_1", - ], }, created_at=AnyStr(), parent_config={ "configurable": { "thread_id": "1", - "checkpoint_ns": AnyStr(), + "checkpoint_ns": "", "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), } }, tasks=( - PregelTask( - id=AnyStr(), - name="grandchild_1", - path=(PULL, "grandchild_1"), - result={"my_key": "hi my value here"}, - ), + PregelTask(id=AnyStr(), name="parent_1", path=(PULL, "parent_1")), ), interrupts=(), ), @@ -6645,70 +6257,416 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: config={ "configurable": { "thread_id": "1", - "checkpoint_ns": AnyStr(), + "checkpoint_ns": "", "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), } }, metadata={ + "parents": {}, "source": "input", - "writes": {"__start__": {"my_key": "hi my value"}}, + "writes": {"my_key": "my value"}, "step": -1, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), "thread_id": "1", - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": [ - "branch:to:child_1", - ], }, created_at=AnyStr(), parent_config=None, tasks=( - PregelTask( - id=AnyStr(), - name="__start__", - path=(PULL, "__start__"), - result={"my_key": "hi my value"}, - ), + PregelTask(id=AnyStr(), name="__start__", path=(PULL, "__start__")), ), interrupts=(), ), - ] - - # replay grandchild checkpoint - assert [ - c - async for c in app.astream( - None, grandchild_history[2].config, subgraphs=True - ) - ] == [ - ( - (AnyStr("child:"), AnyStr("child_1:")), - {"grandchild_1": {"my_key": "hi my value here"}}, + ][0] + ) + # get child graph history + child_history = [ + c async for c in app.aget_state_history(outer_history[2].tasks[0].state) + ] + assert child_history == [ + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "source": "loop", + "writes": {"child_1": {"my_key": "hi my value here and there"}}, + "step": 1, + "parents": {"": AnyStr()}, + "thread_id": "1", + "langgraph_node": "child", + "langgraph_path": [PULL, AnyStr("child")], + "langgraph_step": 2, + "langgraph_triggers": ["branch:to:child"], + "langgraph_checkpoint_ns": AnyStr("child:"), + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + tasks=(), + interrupts=(), + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("child_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "source": "loop", + "writes": None, + "step": 0, + "parents": {"": AnyStr()}, + "thread_id": "1", + "langgraph_node": "child", + "langgraph_path": [PULL, AnyStr("child")], + "langgraph_step": 2, + "langgraph_triggers": ["branch:to:child"], + "langgraph_checkpoint_ns": AnyStr("child:"), + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="child_1", + path=(PULL, "child_1"), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + } + }, + result={"my_key": "hi my value here and there"}, + ), ), - ((), {"__interrupt__": ()}), - ] + interrupts=(), + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), + } + }, + metadata={ + "source": "input", + "writes": {"__start__": {"my_key": "hi my value"}}, + "step": -1, + "parents": {"": AnyStr()}, + "thread_id": "1", + "langgraph_node": "child", + "langgraph_path": [PULL, AnyStr("child")], + "langgraph_step": 2, + "langgraph_triggers": ["branch:to:child"], + "langgraph_checkpoint_ns": AnyStr("child:"), + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=(PULL, "__start__"), + result={"my_key": "hi my value"}, + ), + ), + interrupts=(), + ), + ] + # get grandchild graph history + grandchild_history = [ + c async for c in app.aget_state_history(child_history[1].tasks[0].state) + ] + assert grandchild_history == [ + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), + } + }, + metadata={ + "source": "loop", + "writes": {"grandchild_2": {"my_key": "hi my value here and there"}}, + "step": 2, + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + "thread_id": "1", + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [ + "branch:to:child_1", + ], + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), + } + }, + tasks=(), + interrupts=(), + ), + StateSnapshot( + values={"my_key": "hi my value here"}, + next=("grandchild_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), + } + }, + metadata={ + "source": "loop", + "writes": {"grandchild_1": {"my_key": "hi my value here"}}, + "step": 1, + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + "thread_id": "1", + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [ + "branch:to:child_1", + ], + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="grandchild_2", + path=(PULL, "grandchild_2"), + result={"my_key": "hi my value here and there"}, + ), + ), + interrupts=(), + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("grandchild_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), + } + }, + metadata={ + "source": "loop", + "writes": None, + "step": 0, + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + "thread_id": "1", + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [ + "branch:to:child_1", + ], + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="grandchild_1", + path=(PULL, "grandchild_1"), + result={"my_key": "hi my value here"}, + ), + ), + interrupts=(), + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr(), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), + } + }, + metadata={ + "source": "input", + "writes": {"__start__": {"my_key": "hi my value"}}, + "step": -1, + "parents": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + } + ), + "thread_id": "1", + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [ + "branch:to:child_1", + ], + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=(PULL, "__start__"), + result={"my_key": "hi my value"}, + ), + ), + interrupts=(), + ), + ] + + # replay grandchild checkpoint + assert [ + c async for c in app.astream(None, grandchild_history[2].config, subgraphs=True) + ] == [ + ( + (AnyStr("child:"), AnyStr("child_1:")), + {"grandchild_1": {"my_key": "hi my value here"}}, + ), + ((), {"__interrupt__": ()}), + ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_send_to_nested_graphs(checkpointer_name: str) -> None: +async def test_send_to_nested_graphs(async_checkpointer: BaseCheckpointSaver) -> None: class OverallState(TypedDict): subjects: list[str] jokes: Annotated[list[str], operator.add] @@ -6742,43 +6700,41 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: builder.add_conditional_edges(START, continue_to_jokes) builder.add_edge("generate_joke", END) - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) - config = {"configurable": {"thread_id": "1"}} - tracer = FakeTracer() + graph = builder.compile(checkpointer=async_checkpointer) + config = {"configurable": {"thread_id": "1"}} + tracer = FakeTracer() - # invoke and pause at nested interrupt - assert await graph.ainvoke( - {"subjects": ["cats", "dogs"]}, - config={**config, "callbacks": [tracer]}, - ) == { - "subjects": ["cats", "dogs"], - "jokes": [], - } - assert len(tracer.runs) == 1, "Should produce exactly 1 root run" + # invoke and pause at nested interrupt + assert await graph.ainvoke( + {"subjects": ["cats", "dogs"]}, + config={**config, "callbacks": [tracer]}, + ) == { + "subjects": ["cats", "dogs"], + "jokes": [], + } + assert len(tracer.runs) == 1, "Should produce exactly 1 root run" - # check state - outer_state = await graph.aget_state(config) + # check state + outer_state = await graph.aget_state(config) - # update state of dogs joke graph - await graph.aupdate_state( - outer_state.tasks[1].state, {"subject": "turtles - hohoho"} - ) + # update state of dogs joke graph + await graph.aupdate_state( + outer_state.tasks[1].state, {"subject": "turtles - hohoho"} + ) - # continue past interrupt - assert await graph.ainvoke(None, config=config) == { - "subjects": ["cats", "dogs"], - "jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"], - } + # continue past interrupt + assert await graph.ainvoke(None, config=config) == { + "subjects": ["cats", "dogs"], + "jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"], + } @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_weather_subgraph( - checkpointer_name: str, snapshot: SnapshotAssertion + async_checkpointer: BaseCheckpointSaver, ) -> None: from langchain_core.language_models.fake_chat_models import ( FakeMessagesListChatModel, @@ -6836,9 +6792,6 @@ async def test_weather_subgraph( class RouterState(MessagesState): route: Literal["weather", "other"] - class Router(TypedDict): - route: Literal["weather", "other"] - router_model = FakeMessagesListChatModel( responses=[ AIMessage( @@ -6891,177 +6844,199 @@ async def test_weather_subgraph( def get_first_in_list(): return [*graph.get_state_history(config, limit=1)][0] - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = graph.compile(checkpointer=checkpointer) + graph = graph.compile(checkpointer=async_checkpointer) - config = {"configurable": {"thread_id": "1"}} - thread2 = {"configurable": {"thread_id": "2"}} - inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]} + config = {"configurable": {"thread_id": "1"}} + thread2 = {"configurable": {"thread_id": "2"}} + inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]} - # run with custom output - assert [ - c async for c in graph.astream(inputs, thread2, stream_mode="custom") - ] == [ - "I'm", - " very", - ] - assert [ - c async for c in graph.astream(None, thread2, stream_mode="custom") - ] == [ - " good", - ] + # run with custom output + assert [c async for c in graph.astream(inputs, thread2, stream_mode="custom")] == [ + "I'm", + " very", + ] + assert [c async for c in graph.astream(None, thread2, stream_mode="custom")] == [ + " good", + ] - # run until interrupt - assert [ - c - async for c in graph.astream( - inputs, config=config, stream_mode="updates", subgraphs=True - ) - ] == [ - ((), {"router_node": {"route": "weather"}}), - ((AnyStr("weather_graph:"),), {"model_node": {"city": "San Francisco"}}), - ((), {"__interrupt__": ()}), - ] + # run until interrupt + assert [ + c + async for c in graph.astream( + inputs, config=config, stream_mode="updates", subgraphs=True + ) + ] == [ + ((), {"router_node": {"route": "weather"}}), + ((AnyStr("weather_graph:"),), {"model_node": {"city": "San Francisco"}}), + ((), {"__interrupt__": ()}), + ] - # check current state - state = await graph.aget_state(config) - assert state == StateSnapshot( - values={ - "messages": [_AnyIdHumanMessage(content="what's the weather in sf")], - "route": "weather", - }, - next=("weather_graph",), - config={ + # check current state + state = await graph.aget_state(config) + assert state == StateSnapshot( + values={ + "messages": [_AnyIdHumanMessage(content="what's the weather in sf")], + "route": "weather", + }, + next=("weather_graph",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"router_node": {"route": "weather"}}, + "step": 1, + "parents": {}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": "", "checkpoint_id": AnyStr(), } - }, - metadata={ - "source": "loop", - "writes": {"router_node": {"route": "weather"}}, - "step": 1, - "parents": {}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=({ + } + ), + tasks=( + PregelTask( + id=AnyStr(), + name="weather_graph", + path=(PULL, "weather_graph"), + state={ "configurable": { "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), + "checkpoint_ns": AnyStr("weather_graph:"), } - } + }, ), - tasks=( - PregelTask( - id=AnyStr(), - name="weather_graph", - path=(PULL, "weather_graph"), - state={ + ), + interrupts=(), + ) + # confirm that list() delegates to alist() correctly + assert await asyncio.to_thread(get_first_in_list) == state + + # update + await graph.aupdate_state(state.tasks[0].state, {"city": "la"}) + + # run after update + assert [ + c + async for c in graph.astream( + None, config=config, stream_mode="updates", subgraphs=True + ) + ] == [ + ( + (AnyStr("weather_graph:"),), + { + "weather_node": { + "messages": [{"role": "assistant", "content": "I'ts sunny in la!"}] + } + }, + ), + ( + (), + { + "weather_graph": { + "messages": [ + _AnyIdHumanMessage(content="what's the weather in sf"), + _AnyIdAIMessage(content="I'ts sunny in la!"), + ] + } + }, + ), + ] + + # try updating acting as weather node + config = {"configurable": {"thread_id": "14"}} + inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]} + assert [ + c + async for c in graph.astream( + inputs, config=config, stream_mode="updates", subgraphs=True + ) + ] == [ + ((), {"router_node": {"route": "weather"}}), + ((AnyStr("weather_graph:"),), {"model_node": {"city": "San Francisco"}}), + ((), {"__interrupt__": ()}), + ] + state = await graph.aget_state(config, subgraphs=True) + assert state == StateSnapshot( + values={ + "messages": [_AnyIdHumanMessage(content="what's the weather in sf")], + "route": "weather", + }, + next=("weather_graph",), + config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"router_node": {"route": "weather"}}, + "step": 1, + "parents": {}, + "thread_id": "14", + }, + created_at=AnyStr(), + parent_config=( + { + "configurable": { + "thread_id": "14", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + } + ), + tasks=( + PregelTask( + id=AnyStr(), + name="weather_graph", + path=(PULL, "weather_graph"), + state=StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="what's the weather in sf") + ], + "city": "San Francisco", + }, + next=("weather_node",), + config={ "configurable": { - "thread_id": "1", + "thread_id": "14", "checkpoint_ns": AnyStr("weather_graph:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("weather_graph:"): AnyStr(), + } + ), } }, - ), - ), - interrupts=(), - ) - # confirm that list() delegates to alist() correctly - assert await asyncio.to_thread(get_first_in_list) == state - - # update - await graph.aupdate_state(state.tasks[0].state, {"city": "la"}) - - # run after update - assert [ - c - async for c in graph.astream( - None, config=config, stream_mode="updates", subgraphs=True - ) - ] == [ - ( - (AnyStr("weather_graph:"),), - { - "weather_node": { - "messages": [ - {"role": "assistant", "content": "I'ts sunny in la!"} - ] - } - }, - ), - ( - (), - { - "weather_graph": { - "messages": [ - _AnyIdHumanMessage(content="what's the weather in sf"), - _AnyIdAIMessage(content="I'ts sunny in la!"), - ] - } - }, - ), - ] - - # try updating acting as weather node - config = {"configurable": {"thread_id": "14"}} - inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]} - assert [ - c - async for c in graph.astream( - inputs, config=config, stream_mode="updates", subgraphs=True - ) - ] == [ - ((), {"router_node": {"route": "weather"}}), - ((AnyStr("weather_graph:"),), {"model_node": {"city": "San Francisco"}}), - ((), {"__interrupt__": ()}), - ] - state = await graph.aget_state(config, subgraphs=True) - assert state == StateSnapshot( - values={ - "messages": [_AnyIdHumanMessage(content="what's the weather in sf")], - "route": "weather", - }, - next=("weather_graph",), - config={ - "configurable": { - "thread_id": "14", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"router_node": {"route": "weather"}}, - "step": 1, - "parents": {}, - "thread_id": "14", - }, - created_at=AnyStr(), - parent_config=({ - "configurable": { + metadata={ + "source": "loop", + "writes": {"model_node": {"city": "San Francisco"}}, + "step": 1, + "parents": {"": AnyStr()}, "thread_id": "14", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - tasks=( - PregelTask( - id=AnyStr(), - name="weather_graph", - path=(PULL, "weather_graph"), - state=StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what's the weather in sf") - ], - "city": "San Francisco", - }, - next=("weather_node",), - config={ + "langgraph_node": "weather_graph", + "langgraph_path": [PULL, "weather_graph"], + "langgraph_step": 2, + "langgraph_triggers": ["branch:to:weather_graph"], + "langgraph_checkpoint_ns": AnyStr("weather_graph:"), + }, + created_at=AnyStr(), + parent_config=( + { "configurable": { "thread_id": "14", "checkpoint_ns": AnyStr("weather_graph:"), @@ -7073,166 +7048,140 @@ async def test_weather_subgraph( } ), } - }, - metadata={ - "source": "loop", - "writes": {"model_node": {"city": "San Francisco"}}, - "step": 1, - "parents": {"": AnyStr()}, - "thread_id": "14", - "langgraph_node": "weather_graph", - "langgraph_path": [PULL, "weather_graph"], - "langgraph_step": 2, - "langgraph_triggers": ["branch:to:weather_graph"], - "langgraph_checkpoint_ns": AnyStr("weather_graph:"), - }, - created_at=AnyStr(), - parent_config=({ - "configurable": { - "thread_id": "14", - "checkpoint_ns": AnyStr("weather_graph:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("weather_graph:"): AnyStr(), - } - ), - } - } - ), - interrupts=(), - tasks=( - PregelTask( - id=AnyStr(), - name="weather_node", - path=(PULL, "weather_node"), - ), + } + ), + interrupts=(), + tasks=( + PregelTask( + id=AnyStr(), + name="weather_node", + path=(PULL, "weather_node"), ), ), ), ), - interrupts=(), - ) - await graph.aupdate_state( - state.tasks[0].state.config, - {"messages": [{"role": "assistant", "content": "rainy"}]}, - as_node="weather_node", - ) - state = await graph.aget_state(config, subgraphs=True) - assert state == StateSnapshot( - values={ - "messages": [_AnyIdHumanMessage(content="what's the weather in sf")], - "route": "weather", - }, - next=("weather_graph",), - config={ + ), + interrupts=(), + ) + await graph.aupdate_state( + state.tasks[0].state.config, + {"messages": [{"role": "assistant", "content": "rainy"}]}, + as_node="weather_node", + ) + state = await graph.aget_state(config, subgraphs=True) + assert state == StateSnapshot( + values={ + "messages": [_AnyIdHumanMessage(content="what's the weather in sf")], + "route": "weather", + }, + next=("weather_graph",), + config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"router_node": {"route": "weather"}}, + "step": 1, + "parents": {}, + "thread_id": "14", + }, + created_at=AnyStr(), + parent_config=( + { "configurable": { "thread_id": "14", "checkpoint_ns": "", "checkpoint_id": AnyStr(), } - }, - metadata={ - "source": "loop", - "writes": {"router_node": {"route": "weather"}}, - "step": 1, - "parents": {}, - "thread_id": "14", - }, - created_at=AnyStr(), - parent_config=({ - "configurable": { - "thread_id": "14", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - interrupts=(), - tasks=( - PregelTask( - id=AnyStr(), - name="weather_graph", - path=(PULL, "weather_graph"), - state=StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="what's the weather in sf"), - _AnyIdAIMessage(content="rainy"), - ], - "city": "San Francisco", - }, - next=(), - config={ - "configurable": { - "thread_id": "14", - "checkpoint_ns": AnyStr("weather_graph:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("weather_graph:"): AnyStr(), - } - ), - } - }, - metadata={ - "step": 2, - "source": "update", - "writes": { - "weather_node": { - "messages": [ - {"role": "assistant", "content": "rainy"} - ] - } - }, - "parents": {"": AnyStr()}, - "thread_id": "14", - "checkpoint_id": AnyStr(), - "checkpoint_ns": AnyStr("weather_graph:"), - "langgraph_node": "weather_graph", - "langgraph_path": [PULL, "weather_graph"], - "langgraph_step": 2, - "langgraph_triggers": ["branch:to:weather_graph"], - "langgraph_checkpoint_ns": AnyStr("weather_graph:"), - }, - created_at=AnyStr(), - parent_config=({ - "configurable": { - "thread_id": "14", - "checkpoint_ns": AnyStr("weather_graph:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("weather_graph:"): AnyStr(), - } - ), - } - } - ), - tasks=(), - interrupts=(), - ), - ), - ), - ) - assert [ - c - async for c in graph.astream( - None, config=config, stream_mode="updates", subgraphs=True - ) - ] == [ - ( - (), - { - "weather_graph": { + } + ), + interrupts=(), + tasks=( + PregelTask( + id=AnyStr(), + name="weather_graph", + path=(PULL, "weather_graph"), + state=StateSnapshot( + values={ "messages": [ _AnyIdHumanMessage(content="what's the weather in sf"), _AnyIdAIMessage(content="rainy"), - ] - } - }, + ], + "city": "San Francisco", + }, + next=(), + config={ + "configurable": { + "thread_id": "14", + "checkpoint_ns": AnyStr("weather_graph:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("weather_graph:"): AnyStr(), + } + ), + } + }, + metadata={ + "step": 2, + "source": "update", + "writes": { + "weather_node": { + "messages": [{"role": "assistant", "content": "rainy"}] + } + }, + "parents": {"": AnyStr()}, + "thread_id": "14", + "checkpoint_id": AnyStr(), + "checkpoint_ns": AnyStr("weather_graph:"), + "langgraph_node": "weather_graph", + "langgraph_path": [PULL, "weather_graph"], + "langgraph_step": 2, + "langgraph_triggers": ["branch:to:weather_graph"], + "langgraph_checkpoint_ns": AnyStr("weather_graph:"), + }, + created_at=AnyStr(), + parent_config=( + { + "configurable": { + "thread_id": "14", + "checkpoint_ns": AnyStr("weather_graph:"), + "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("weather_graph:"): AnyStr(), + } + ), + } + } + ), + tasks=(), + interrupts=(), + ), ), - ] + ), + ) + assert [ + c + async for c in graph.astream( + None, config=config, stream_mode="updates", subgraphs=True + ) + ] == [ + ( + (), + { + "weather_graph": { + "messages": [ + _AnyIdHumanMessage(content="what's the weather in sf"), + _AnyIdAIMessage(content="rainy"), + ] + } + }, + ), + ] diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index cc104c288..fe85ba99f 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -66,7 +66,6 @@ from langgraph.types import ( ) from tests.agents import AgentAction, AgentFinish from tests.any_str import AnyStr, AnyVersion, FloatBetween, UnsortedSequence -from tests.conftest import ALL_CHECKPOINTERS_SYNC, SHOULD_CHECK_SNAPSHOTS from tests.messages import ( _AnyIdAIMessage, _AnyIdAIMessageChunk, @@ -491,22 +490,21 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: graph.set_finish_point("add_one") gapp = graph.compile() - if SHOULD_CHECK_SNAPSHOTS: - assert app.input_schema.model_json_schema() == { - "title": "LangGraphInput", - "type": "integer", + assert app.input_schema.model_json_schema() == { + "title": "LangGraphInput", + "type": "integer", + } + assert app.output_schema.model_json_schema() == { + "title": "LangGraphOutput", + "type": "integer", + } + with warnings.catch_warnings(): + warnings.simplefilter("error") # raise warnings as errors + assert app.config_schema().model_json_schema() == { + "properties": {}, + "title": "LangGraphConfig", + "type": "object", } - assert app.output_schema.model_json_schema() == { - "title": "LangGraphOutput", - "type": "integer", - } - with warnings.catch_warnings(): - warnings.simplefilter("error") # raise warnings as errors - assert app.config_schema().model_json_schema() == { - "properties": {}, - "title": "LangGraphConfig", - "type": "object", - } assert app.invoke(2) == 3 assert app.invoke(2, output_keys=["output"]) == {"output": 3} @@ -548,24 +546,23 @@ def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None: input_channels="input", ) - if SHOULD_CHECK_SNAPSHOTS: - assert app.input_schema.model_json_schema() == { - "title": "LangGraphInput", - "type": "integer", - } - assert app.output_schema.model_json_schema() == { - "title": "LangGraphOutput", - "type": "object", - "properties": { - "output": {"title": "Output", "type": "integer", "default": None}, - "fixed": {"title": "Fixed", "type": "integer", "default": None}, - "output_plus_one": { - "title": "Output Plus One", - "type": "integer", - "default": None, - }, + assert app.input_schema.model_json_schema() == { + "title": "LangGraphInput", + "type": "integer", + } + assert app.output_schema.model_json_schema() == { + "title": "LangGraphOutput", + "type": "object", + "properties": { + "output": {"title": "Output", "type": "integer", "default": None}, + "fixed": {"title": "Fixed", "type": "integer", "default": None}, + "output_plus_one": { + "title": "Output Plus One", + "type": "integer", + "default": None, }, - } + }, + } assert app.invoke(2) == {"output": 3, "fixed": 5, "output_plus_one": 4} @@ -580,18 +577,17 @@ def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: output_channels=["output"], ) - if SHOULD_CHECK_SNAPSHOTS: - assert app.input_schema.model_json_schema() == { - "title": "LangGraphInput", - "type": "integer", - } - assert app.output_schema.model_json_schema() == { - "title": "LangGraphOutput", - "type": "object", - "properties": { - "output": {"title": "Output", "type": "integer", "default": None} - }, - } + assert app.input_schema.model_json_schema() == { + "title": "LangGraphInput", + "type": "integer", + } + assert app.output_schema.model_json_schema() == { + "title": "LangGraphOutput", + "type": "object", + "properties": { + "output": {"title": "Output", "type": "integer", "default": None} + }, + } assert app.invoke(2) == {"output": 3} @@ -605,21 +601,18 @@ def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None: input_channels=["input"], output_channels=["output"], ) - if SHOULD_CHECK_SNAPSHOTS: - assert app.input_schema.model_json_schema() == { - "title": "LangGraphInput", - "type": "object", - "properties": { - "input": {"title": "Input", "type": "integer", "default": None} - }, - } - assert app.output_schema.model_json_schema() == { - "title": "LangGraphOutput", - "type": "object", - "properties": { - "output": {"title": "Output", "type": "integer", "default": None} - }, - } + assert app.input_schema.model_json_schema() == { + "title": "LangGraphInput", + "type": "object", + "properties": {"input": {"title": "Input", "type": "integer", "default": None}}, + } + assert app.output_schema.model_json_schema() == { + "title": "LangGraphOutput", + "type": "object", + "properties": { + "output": {"title": "Output", "type": "integer", "default": None} + }, + } assert app.invoke({"input": 2}) == {"output": 3} @@ -1022,13 +1015,9 @@ def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> Non assert app.invoke(2) == [3, 3] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_invoke_checkpoint_two( - mocker: MockerFixture, request: pytest.FixtureRequest, checkpointer_name: str + mocker: MockerFixture, sync_checkpointer: BaseCheckpointSaver ) -> None: - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) errored_once = False @@ -1060,26 +1049,26 @@ def test_invoke_checkpoint_two( }, input_channels="input", output_channels="output", - checkpointer=checkpointer, + checkpointer=sync_checkpointer, retry_policy=RetryPolicy(), ) # total starts out as 0, so output is 0+2=2 assert app.invoke(2, {"configurable": {"thread_id": "1"}}) == 2 - checkpoint = checkpointer.get({"configurable": {"thread_id": "1"}}) + checkpoint = sync_checkpointer.get({"configurable": {"thread_id": "1"}}) assert checkpoint is not None assert checkpoint["channel_values"].get("total") == 2 # total is now 2, so output is 2+3=5 assert app.invoke(3, {"configurable": {"thread_id": "1"}}) == 5 assert errored_once, "errored and retried" - checkpoint_tup = checkpointer.get_tuple({"configurable": {"thread_id": "1"}}) + checkpoint_tup = sync_checkpointer.get_tuple({"configurable": {"thread_id": "1"}}) assert checkpoint_tup is not None assert checkpoint_tup.checkpoint["channel_values"].get("total") == 7 # total is now 2+5=7, so output would be 7+4=11, but raises ValueError with pytest.raises(ValueError): app.invoke(4, {"configurable": {"thread_id": "1"}}) # checkpoint is not updated, error is recorded - checkpoint_tup = checkpointer.get_tuple({"configurable": {"thread_id": "1"}}) + checkpoint_tup = sync_checkpointer.get_tuple({"configurable": {"thread_id": "1"}}) assert checkpoint_tup is not None assert checkpoint_tup.checkpoint["channel_values"].get("total") == 7 assert checkpoint_tup.pending_writes == [ @@ -1087,23 +1076,17 @@ def test_invoke_checkpoint_two( ] # on a new thread, total starts out as 0, so output is 0+5=5 assert app.invoke(5, {"configurable": {"thread_id": "2"}}) == 5 - checkpoint = checkpointer.get({"configurable": {"thread_id": "1"}}) + checkpoint = sync_checkpointer.get({"configurable": {"thread_id": "1"}}) assert checkpoint is not None assert checkpoint["channel_values"].get("total") == 7 - checkpoint = checkpointer.get({"configurable": {"thread_id": "2"}}) + checkpoint = sync_checkpointer.get({"configurable": {"thread_id": "2"}}) assert checkpoint is not None assert checkpoint["channel_values"].get("total") == 5 -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_pending_writes_resume( - request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool + sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool ) -> None: - - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - class State(TypedDict): value: Annotated[int, operator.add] @@ -1133,7 +1116,7 @@ def test_pending_writes_resume( ) builder.add_edge(START, "one") builder.add_edge(START, "two") - graph = builder.compile(checkpointer=checkpointer) + graph = builder.compile(checkpointer=sync_checkpointer) thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} with pytest.raises(ConnectionError, match="I'm not good"): @@ -1166,7 +1149,7 @@ def test_pending_writes_resume( assert state.values == {"value": 1} assert state.next == ("one", "two") # should contain pending write of "one" - checkpoint = checkpointer.get_tuple(thread1) + checkpoint = sync_checkpointer.get_tuple(thread1) assert checkpoint is not None # should contain error from "two" expected_writes = [ @@ -1201,9 +1184,8 @@ def test_pending_writes_resume( "value": 6 } - # check all final checkpoints - checkpoints = [c for c in checkpointer.list(thread1)] + checkpoints = [c for c in sync_checkpointer.list(thread1)] # we should have 3 assert len(checkpoints) == (3 if checkpoint_during else 2) # the last one not too interesting for this test @@ -1488,12 +1470,9 @@ def test_send_sequences() -> None: ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_imp_task( - request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool + sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool ) -> None: - - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") mapper_calls = 0 class Configurable: @@ -1506,7 +1485,7 @@ def test_imp_task( time.sleep(input / 100) return str(input) * 2 - @entrypoint(checkpointer=checkpointer, config_schema=Configurable) + @entrypoint(checkpointer=sync_checkpointer, config_schema=Configurable) def graph(input: list[int]) -> list[str]: futures = [mapper(i) for i in input] mapped = [f.result() for f in futures] @@ -1583,13 +1562,9 @@ def test_imp_task( assert mapper_calls == 2 -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_imp_nested( - request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool + sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool ) -> None: - - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - def mynode(input: list[str]) -> list[str]: return [it + "a" for it in input] @@ -1609,7 +1584,7 @@ def test_imp_nested( time.sleep(input / 100) return sub.result() * 2 - @entrypoint(checkpointer=checkpointer) + @entrypoint(checkpointer=sync_checkpointer) def graph(input: list[int]) -> list[str]: futures = [mapper(i) for i in input] mapped = [f.result() for f in futures] @@ -1654,13 +1629,9 @@ def test_imp_nested( ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_imp_stream_order( - request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool + sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool ) -> None: - - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - @task() def foo(state: dict) -> tuple: return state["a"] + "foo", "bar" @@ -1673,7 +1644,7 @@ def test_imp_stream_order( def baz(state: dict) -> dict: return {"a": state["a"] + "baz", "c": "something else"} - @entrypoint(checkpointer=checkpointer) + @entrypoint(checkpointer=sync_checkpointer) def graph(state: dict) -> dict: fut_foo = foo(state) fut_bar = bar(*fut_foo.result()) @@ -1699,11 +1670,9 @@ def test_imp_stream_order( assert graph.get_state(thread1).values == {"a": "0foobarbaz", "c": "something else"} -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_invoke_checkpoint_three( - mocker: MockerFixture, request: pytest.FixtureRequest, checkpointer_name: str + mocker: MockerFixture, sync_checkpointer: BaseCheckpointSaver ) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") adder = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) def raise_if_above_10(input: int) -> int: @@ -1727,7 +1696,7 @@ def test_invoke_checkpoint_three( }, input_channels="input", output_channels="output", - checkpointer=checkpointer, + checkpointer=sync_checkpointer, ) thread_1 = {"configurable": {"thread_id": "1"}} @@ -1739,7 +1708,7 @@ def test_invoke_checkpoint_three( assert state.next == () assert ( state.config["configurable"]["checkpoint_id"] - == checkpointer.get(thread_1)["id"] + == sync_checkpointer.get(thread_1)["id"] ) # total is now 2, so output is 2+3=5 assert app.invoke(3, thread_1) == 5 @@ -1748,7 +1717,7 @@ def test_invoke_checkpoint_three( assert state.values.get("total") == 7 assert ( state.config["configurable"]["checkpoint_id"] - == checkpointer.get(thread_1)["id"] + == sync_checkpointer.get(thread_1)["id"] ) # total is now 2+5=7, so output would be 7+4=11, but raises ValueError with pytest.raises(ValueError): @@ -1778,7 +1747,6 @@ def test_invoke_checkpoint_three( assert state.values.get("total") == 5 assert state.next == () - assert len(list(app.get_state_history(thread_1, limit=1))) == 1 # list all checkpoints for thread 1 thread_1_history = [c for c in app.get_state_history(thread_1)] @@ -1805,11 +1773,11 @@ def test_invoke_checkpoint_three( assert thread_1_history[-2].values["total"] == 2 # can get each checkpoint using aget with config assert ( - checkpointer.get(thread_1_history[0].config)["id"] + sync_checkpointer.get(thread_1_history[0].config)["id"] == thread_1_history[0].config["configurable"]["checkpoint_id"] ) assert ( - checkpointer.get(thread_1_history[1].config)["id"] + sync_checkpointer.get(thread_1_history[1].config)["id"] == thread_1_history[1].config["configurable"]["checkpoint_id"] ) @@ -1869,12 +1837,9 @@ def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None assert [*executor.map(app.invoke, [2] * 100)] == [[13, 13]] * 100 -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_invoke_join_then_call_other_pregel( - mocker: MockerFixture, request: pytest.FixtureRequest, checkpointer_name: str + mocker: MockerFixture, sync_checkpointer: BaseCheckpointSaver ) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - add_one = mocker.Mock(side_effect=lambda x: x + 1) add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x]) @@ -1926,7 +1891,7 @@ def test_invoke_join_then_call_other_pregel( assert [*executor.map(app.invoke, [[2, 3]] * 10)] == [27] * 10 # add checkpointer - app.checkpointer = checkpointer + app.checkpointer = sync_checkpointer # subgraph is called twice in the same node, but that works assert app.invoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27 @@ -2026,11 +1991,10 @@ def test_conditional_entrypoint_graph(snapshot: SnapshotAssertion) -> None: app = workflow.compile() - if SHOULD_CHECK_SNAPSHOTS: - assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot - assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot + assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot + assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot assert ( app.invoke("what is weather in sf", debug=True) @@ -2068,11 +2032,10 @@ def test_conditional_entrypoint_to_multiple_state_graph( app = workflow.compile() - if SHOULD_CHECK_SNAPSHOTS: - assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot - assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot + assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot + assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot assert app.invoke({"locations": ["sf", "nyc"]}, debug=True) == { "locations": ["sf", "nyc"], @@ -2193,10 +2156,9 @@ def test_state_graph_w_config_inherited_state_keys(snapshot: SnapshotAssertion) app = builder.compile() - if SHOULD_CHECK_SNAPSHOTS: - assert json.dumps(app.config_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot + assert json.dumps(app.config_schema().model_json_schema()) == snapshot + assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot + assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot assert builder.channels.keys() == {"input", "agent_outcome", "intermediate_steps"} @@ -2259,11 +2221,10 @@ def test_conditional_entrypoint_graph_state(snapshot: SnapshotAssertion) -> None app = workflow.compile() - if SHOULD_CHECK_SNAPSHOTS: - assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot - assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot - assert app.get_graph().draw_mermaid(with_styles=False) == snapshot + assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot + assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot + assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot assert app.invoke({"input": "what is weather in sf"}) == { "input": "what is weather in sf", @@ -2276,14 +2237,9 @@ def test_conditional_entrypoint_graph_state(snapshot: SnapshotAssertion) -> None ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_in_one_fan_out_state_graph_waiting_edge( - snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str + snapshot: SnapshotAssertion, sync_checkpointer: BaseCheckpointSaver ) -> None: - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -2331,7 +2287,7 @@ def test_in_one_fan_out_state_graph_waiting_edge( app = workflow.compile() - if checkpointer_name == "memory": + if isinstance(sync_checkpointer, InMemorySaver): assert app.get_graph().draw_mermaid(with_styles=False) == snapshot assert app.invoke({"query": "what is weather in sf"}) == { @@ -2349,7 +2305,7 @@ def test_in_one_fan_out_state_graph_waiting_edge( ] app_w_interrupt = workflow.compile( - checkpointer=checkpointer, + checkpointer=sync_checkpointer, interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -2369,7 +2325,7 @@ def test_in_one_fan_out_state_graph_waiting_edge( ] app_w_interrupt = workflow.compile( - checkpointer=checkpointer, + checkpointer=sync_checkpointer, interrupt_before=["qa"], ) config = {"configurable": {"thread_id": "2"}} @@ -2385,8 +2341,9 @@ def test_in_one_fan_out_state_graph_waiting_edge( ] app_w_interrupt.update_state(config, {"docs": ["doc5"]}) - expected_parent_config = (list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ) + expected_parent_config = list(app_w_interrupt.checkpointer.list(config, limit=2))[ + -1 + ].config assert app_w_interrupt.get_state(config) == StateSnapshot( values={ "query": "analyzed: query: what is weather in sf", @@ -2419,17 +2376,11 @@ def test_in_one_fan_out_state_graph_waiting_edge( @pytest.mark.parametrize("use_waiting_edge", (True, False)) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_in_one_fan_out_state_graph_defer_node( snapshot: SnapshotAssertion, - request: pytest.FixtureRequest, - checkpointer_name: str, + sync_checkpointer: BaseCheckpointSaver, use_waiting_edge: bool, ) -> None: - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -2481,7 +2432,7 @@ def test_in_one_fan_out_state_graph_defer_node( app = workflow.compile() - if checkpointer_name == "memory": + if isinstance(sync_checkpointer, InMemorySaver): assert app.get_graph().draw_mermaid(with_styles=False) == snapshot assert app.invoke({"query": "what is weather in sf"}) == { @@ -2625,7 +2576,7 @@ def test_in_one_fan_out_state_graph_defer_node( ] app_w_interrupt = workflow.compile( - checkpointer=checkpointer, + checkpointer=sync_checkpointer, interrupt_after=["analyzer_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -2645,7 +2596,7 @@ def test_in_one_fan_out_state_graph_defer_node( ] app_w_interrupt = workflow.compile( - checkpointer=checkpointer, + checkpointer=sync_checkpointer, interrupt_before=["qa"], ) config = {"configurable": {"thread_id": "2"}} @@ -2661,8 +2612,9 @@ def test_in_one_fan_out_state_graph_defer_node( ] app_w_interrupt.update_state(config, {"docs": ["doc5"]}) - expected_parent_config = (list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ) + expected_parent_config = list(app_w_interrupt.checkpointer.list(config, limit=2))[ + -1 + ].config assert app_w_interrupt.get_state(config) == StateSnapshot( values={ "query": "analyzed: query: what is weather in sf", @@ -2695,17 +2647,11 @@ def test_in_one_fan_out_state_graph_defer_node( @pytest.mark.parametrize("with_path_map", (True, False)) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_in_one_fan_out_state_graph_then_defer_node( snapshot: SnapshotAssertion, - request: pytest.FixtureRequest, - checkpointer_name: str, + sync_checkpointer: BaseCheckpointSaver, with_path_map: bool, ) -> None: - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -2755,7 +2701,7 @@ def test_in_one_fan_out_state_graph_then_defer_node( app = workflow.compile() - if checkpointer_name == "memory" and with_path_map: + if isinstance(sync_checkpointer, InMemorySaver) and with_path_map: assert app.get_graph().draw_mermaid(with_styles=False) == snapshot assert app.invoke({"query": "what is weather in sf"}) == { @@ -2900,7 +2846,7 @@ def test_in_one_fan_out_state_graph_then_defer_node( ] app_w_interrupt = workflow.compile( - checkpointer=checkpointer, + checkpointer=sync_checkpointer, interrupt_after=["analyzer_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -2920,7 +2866,7 @@ def test_in_one_fan_out_state_graph_then_defer_node( ] app_w_interrupt = workflow.compile( - checkpointer=checkpointer, + checkpointer=sync_checkpointer, interrupt_before=["qa"], ) config = {"configurable": {"thread_id": "2"}} @@ -2936,8 +2882,9 @@ def test_in_one_fan_out_state_graph_then_defer_node( ] app_w_interrupt.update_state(config, {"docs": ["doc5"]}) - expected_parent_config = (list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ) + expected_parent_config = list(app_w_interrupt.checkpointer.list(config, limit=2))[ + -1 + ].config assert app_w_interrupt.get_state(config) == StateSnapshot( values={ "query": "analyzed: query: what is weather in sf", @@ -2969,14 +2916,9 @@ def test_in_one_fan_out_state_graph_then_defer_node( ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_in_one_fan_out_state_graph_waiting_edge_via_branch( - snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str + snapshot: SnapshotAssertion, sync_checkpointer: BaseCheckpointSaver ) -> None: - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -3027,7 +2969,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch( app = workflow.compile() - if checkpointer_name == "memory": + if isinstance(sync_checkpointer, InMemorySaver): assert app.get_graph().draw_mermaid(with_styles=False) == snapshot assert app.invoke({"query": "what is weather in sf"}, debug=True) == { @@ -3045,7 +2987,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch( ] app_w_interrupt = workflow.compile( - checkpointer=checkpointer, + checkpointer=sync_checkpointer, interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -3065,14 +3007,10 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch( ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( snapshot: SnapshotAssertion, - request: pytest.FixtureRequest, - checkpointer_name: str, + sync_checkpointer: BaseCheckpointSaver, ) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -3150,7 +3088,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( app = workflow.compile() - if SHOULD_CHECK_SNAPSHOTS and checkpointer_name == "memory": + if isinstance(sync_checkpointer, InMemorySaver): assert app.get_graph().draw_mermaid(with_styles=False) == snapshot assert app.get_input_schema().model_json_schema() == snapshot assert app.get_output_schema().model_json_schema() == snapshot @@ -3172,7 +3110,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( ] app_w_interrupt = workflow.compile( - checkpointer=checkpointer, + checkpointer=sync_checkpointer, interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -3205,13 +3143,9 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( } -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input( - request: pytest.FixtureRequest, - checkpointer_name: str, + sync_checkpointer: BaseCheckpointSaver, ) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -3301,7 +3235,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_inp ] app_w_interrupt = workflow.compile( - checkpointer=checkpointer, + checkpointer=sync_checkpointer, interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -3334,14 +3268,9 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_inp } -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( - request: pytest.FixtureRequest, checkpointer_name: str + sync_checkpointer: BaseCheckpointSaver, ) -> None: - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -3420,7 +3349,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( ) app_w_interrupt = workflow.compile( - checkpointer=checkpointer, + checkpointer=sync_checkpointer, interrupt_after=["retriever_one"], ) config = {"configurable": {"thread_id": "1"}} @@ -3864,13 +3793,9 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None: ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_subgraph_checkpoint_true( - request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool + sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool ) -> None: - - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - class InnerState(TypedDict): my_key: Annotated[str, operator.add] my_other_key: str @@ -3897,7 +3822,7 @@ def test_subgraph_checkpoint_true( graph.add_conditional_edges( "inner", lambda s: "inner" if s["my_key"].count("there") < 2 else END ) - app = graph.compile(checkpointer=checkpointer) + app = graph.compile(checkpointer=sync_checkpointer) config = {"configurable": {"thread_id": "2"}} assert [ @@ -3930,13 +3855,9 @@ def test_subgraph_checkpoint_true( ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_subgraph_checkpoint_true_interrupt( - request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool + sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool ) -> None: - - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - # Define subgraph class SubgraphState(TypedDict): # note that none of these keys are shared with the parent graph state @@ -3973,7 +3894,7 @@ def test_subgraph_checkpoint_true_interrupt( builder.add_edge(START, "node_1") builder.add_edge("node_1", "node_2") - graph = builder.compile(checkpointer=checkpointer) + graph = builder.compile(checkpointer=sync_checkpointer) config = {"configurable": {"thread_id": "1"}} assert graph.invoke( @@ -3996,12 +3917,9 @@ def test_subgraph_checkpoint_true_interrupt( ) == {"foo": "hi! foobaz"} -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_stream_subgraphs_during_execution( - request: pytest.FixtureRequest, checkpointer_name: str + sync_checkpointer: BaseCheckpointSaver, ) -> None: - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - class InnerState(TypedDict): my_key: Annotated[str, operator.add] my_other_key: str @@ -4043,7 +3961,7 @@ def test_stream_subgraphs_during_execution( graph.add_edge(["inner", "outer_1"], "outer_2") graph.add_edge("outer_2", END) - app = graph.compile(checkpointer=checkpointer) + app = graph.compile(checkpointer=sync_checkpointer) start = time.perf_counter() chunks: list[tuple[float, Any]] = [] @@ -4076,12 +3994,7 @@ def test_stream_subgraphs_during_execution( ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_stream_buffering_single_node( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - +def test_stream_buffering_single_node(sync_checkpointer: BaseCheckpointSaver) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] @@ -4095,7 +4008,7 @@ def test_stream_buffering_single_node( builder.add_node("node", node) builder.add_edge(START, "node") builder.add_edge("node", END) - graph = builder.compile(checkpointer=checkpointer) + graph = builder.compile(checkpointer=sync_checkpointer) start = time.perf_counter() chunks: list[tuple[float, Any]] = [] @@ -4109,13 +4022,9 @@ def test_stream_buffering_single_node( ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_nested_graph_interrupts_parallel( - request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool + sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool ) -> None: - - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - class InnerState(TypedDict): my_key: Annotated[str, operator.add] my_other_key: str @@ -4156,7 +4065,7 @@ def test_nested_graph_interrupts_parallel( graph.add_edge(["inner", "outer_1"], "outer_2") graph.set_finish_point("outer_2") - app = graph.compile(checkpointer=checkpointer) + app = graph.compile(checkpointer=sync_checkpointer) # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} @@ -4213,7 +4122,7 @@ def test_nested_graph_interrupts_parallel( ] # test interrupts BEFORE the parallel node - app = graph.compile(checkpointer=checkpointer, interrupt_before=["outer_1"]) + app = graph.compile(checkpointer=sync_checkpointer, interrupt_before=["outer_1"]) config = {"configurable": {"thread_id": "4"}} assert [ *app.stream( @@ -4243,7 +4152,7 @@ def test_nested_graph_interrupts_parallel( ] # test interrupts AFTER the parallel node - app = graph.compile(checkpointer=checkpointer, interrupt_after=["outer_1"]) + app = graph.compile(checkpointer=sync_checkpointer, interrupt_after=["outer_1"]) config = {"configurable": {"thread_id": "5"}} assert [ *app.stream( @@ -4274,13 +4183,9 @@ def test_nested_graph_interrupts_parallel( ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_doubly_nested_graph_interrupts( - request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool + sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool ) -> None: - - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - class State(TypedDict): my_key: str @@ -4328,7 +4233,7 @@ def test_doubly_nested_graph_interrupts( graph.add_edge("child", "parent_2") graph.set_finish_point("parent_2") - app = graph.compile(checkpointer=checkpointer) + app = graph.compile(checkpointer=sync_checkpointer) # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} @@ -4799,10 +4704,7 @@ def test_xray_lance(snapshot: SnapshotAssertion): assert graph.get_graph(xray=1).to_json() == snapshot -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_channel_values(request: pytest.FixtureRequest, checkpointer_name: str) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - +def test_channel_values(sync_checkpointer: BaseCheckpointSaver) -> None: config = {"configurable": {"thread_id": "1"}} chain = Channel.subscribe_to("input") | Channel.write_to("output") app = Pregel( @@ -4816,10 +4718,10 @@ def test_channel_values(request: pytest.FixtureRequest, checkpointer_name: str) }, input_channels=["input", "ephemeral"], output_channels="output", - checkpointer=checkpointer, + checkpointer=sync_checkpointer, ) app.invoke({"input": 1, "ephemeral": "meow"}, config) - assert checkpointer.get(config)["channel_values"] == {"input": 1, "output": 1} + assert sync_checkpointer.get(config)["channel_values"] == {"input": 1, "output": 1} def test_xray_issue(snapshot: SnapshotAssertion) -> None: @@ -4928,12 +4830,9 @@ def test_multiple_sinks_subgraphs(snapshot: SnapshotAssertion) -> None: assert app.get_graph(xray=True).draw_mermaid() == snapshot -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_store_injected( - request: pytest.FixtureRequest, checkpointer_name: str, sync_store: BaseStore + sync_checkpointer: BaseCheckpointSaver, sync_store: BaseStore ) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - class State(TypedDict): count: Annotated[int, operator.add] @@ -4976,7 +4875,7 @@ def test_store_injected( builder.add_node(f"node_{i}", Node(i)) builder.add_edge("__start__", f"node_{i}") - graph = builder.compile(store=sync_store, checkpointer=checkpointer) + graph = builder.compile(store=sync_store, checkpointer=sync_checkpointer) results = graph.batch( [{"count": 0}] * M, @@ -5403,8 +5302,7 @@ def test_runnable_passthrough_node_graph() -> None: assert graph.get_graph(xray=True).to_json() == graph.get_graph(xray=False).to_json() -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_parent_command(request: pytest.FixtureRequest, checkpointer_name: str) -> None: +def test_parent_command(sync_checkpointer: BaseCheckpointSaver) -> None: from langchain_core.messages import BaseMessage from langchain_core.tools import tool @@ -5426,8 +5324,7 @@ def test_parent_command(request: pytest.FixtureRequest, checkpointer_name: str) builder = StateGraph(CustomParentState) builder.add_node("alice", subgraph) builder.add_edge(START, "alice") - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - graph = builder.compile(checkpointer=checkpointer) + graph = builder.compile(checkpointer=sync_checkpointer) config = {"configurable": {"thread_id": "1"}} @@ -5468,7 +5365,8 @@ def test_parent_command(request: pytest.FixtureRequest, checkpointer_name: str) "parents": {}, }, created_at=AnyStr(), - parent_config=({ + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": "", @@ -5481,10 +5379,7 @@ def test_parent_command(request: pytest.FixtureRequest, checkpointer_name: str) ) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_interrupt_subgraph(request: pytest.FixtureRequest, checkpointer_name: str): - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - +def test_interrupt_subgraph(sync_checkpointer: BaseCheckpointSaver): class State(TypedDict): baz: str @@ -5504,7 +5399,7 @@ def test_interrupt_subgraph(request: pytest.FixtureRequest, checkpointer_name: s builder.add_node("bar", child_builder.compile()) builder.add_edge(START, "foo") builder.add_edge("foo", "bar") - graph = builder.compile(checkpointer=checkpointer) + graph = builder.compile(checkpointer=sync_checkpointer) thread1 = {"configurable": {"thread_id": "1"}} # First run, interrupted at bar @@ -5513,10 +5408,7 @@ def test_interrupt_subgraph(request: pytest.FixtureRequest, checkpointer_name: s assert graph.invoke(Command(resume="bar"), thread1) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_interrupt_multiple(request: pytest.FixtureRequest, checkpointer_name: str): - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - +def test_interrupt_multiple(sync_checkpointer: BaseCheckpointSaver): class State(TypedDict): my_key: Annotated[str, operator.add] @@ -5529,7 +5421,7 @@ def test_interrupt_multiple(request: pytest.FixtureRequest, checkpointer_name: s builder.add_node("node", node) builder.add_edge(START, "node") - graph = builder.compile(checkpointer=checkpointer) + graph = builder.compile(checkpointer=sync_checkpointer) thread1 = {"configurable": {"thread_id": "1"}} assert [e for e in graph.stream({"my_key": "DE", "market": "DE"}, thread1)] == [ @@ -5568,10 +5460,7 @@ def test_interrupt_multiple(request: pytest.FixtureRequest, checkpointer_name: s ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_interrupt_loop(request: pytest.FixtureRequest, checkpointer_name: str): - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - +def test_interrupt_loop(sync_checkpointer: BaseCheckpointSaver): class State(TypedDict): age: int other: str @@ -5594,7 +5483,7 @@ def test_interrupt_loop(request: pytest.FixtureRequest, checkpointer_name: str): builder.add_node("node", ask_age) builder.add_edge(START, "node") - graph = builder.compile(checkpointer=checkpointer) + graph = builder.compile(checkpointer=sync_checkpointer) thread1 = {"configurable": {"thread_id": "1"}} assert [e for e in graph.stream({"other": ""}, thread1)] == [ @@ -5653,14 +5542,9 @@ def test_interrupt_loop(request: pytest.FixtureRequest, checkpointer_name: str): ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_interrupt_functional( - request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion + sync_checkpointer: BaseCheckpointSaver, snapshot: SnapshotAssertion ) -> None: - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - @task def foo(state: dict) -> dict: return {"a": state["a"] + "foo"} @@ -5669,7 +5553,7 @@ def test_interrupt_functional( def bar(state: dict) -> dict: return {"a": state["a"] + "bar", "b": state["b"]} - @entrypoint(checkpointer=checkpointer) + @entrypoint(checkpointer=sync_checkpointer) def graph(inputs: dict) -> dict: fut_foo = foo(inputs) value = interrupt("Provide value for bar:") @@ -5693,14 +5577,9 @@ def test_interrupt_functional( assert res == {"a": "foobar", "b": "bar"} -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_interrupt_task_functional( - request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion + sync_checkpointer: BaseCheckpointSaver, snapshot: SnapshotAssertion ) -> None: - checkpointer: BaseCheckpointSaver = request.getfixturevalue( - f"checkpointer_{checkpointer_name}" - ) - @task def foo(state: dict) -> dict: return {"a": state["a"] + "foo"} @@ -5710,7 +5589,7 @@ def test_interrupt_task_functional( value = interrupt("Provide value for bar:") return {"a": state["a"] + value} - @entrypoint(checkpointer=checkpointer) + @entrypoint(checkpointer=sync_checkpointer) def graph(inputs: dict) -> dict: fut_foo = foo(inputs) fut_bar = bar(fut_foo.result()) @@ -5734,7 +5613,7 @@ def test_interrupt_task_functional( # Test that we can interrupt the same task multiple times config = {"configurable": {"thread_id": "2"}} - @entrypoint(checkpointer=checkpointer) + @entrypoint(checkpointer=sync_checkpointer) def graph(inputs: dict) -> dict: foo_result = foo(inputs).result() bar_result = bar(foo_result).result() @@ -5811,14 +5690,11 @@ def test_command_pydantic_dataclass() -> None: assert graph.invoke(State(foo="")) == {"foo": "foobar"} -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_command_with_static_breakpoints( - request: pytest.FixtureRequest, checkpointer_name: str + sync_checkpointer: BaseCheckpointSaver, ) -> None: """Test that we can use Command to resume and update with static breakpoints.""" - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - class State(TypedDict): """The graph state.""" @@ -5840,7 +5716,7 @@ def test_command_with_static_breakpoints( builder.add_edge(START, "node1") builder.add_edge("node1", "node2") - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["node1"]) + graph = builder.compile(checkpointer=sync_checkpointer, interrupt_before=["node1"]) config = {"configurable": {"thread_id": str(uuid.uuid4())}} # Start the graph and interrupt at the first node @@ -5849,12 +5725,9 @@ def test_command_with_static_breakpoints( assert result == {"foo": "abc|node-1|node-2"} -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_multistep_plan(request: pytest.FixtureRequest, checkpointer_name: str): +def test_multistep_plan(sync_checkpointer: BaseCheckpointSaver): from langchain_core.messages import AnyMessage - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - class State(TypedDict, total=False): plan: list[Union[str, list[str]]] messages: Annotated[list[AnyMessage], add_messages] @@ -5894,7 +5767,7 @@ def test_multistep_plan(request: pytest.FixtureRequest, checkpointer_name: str): builder.add_node(step3) builder.add_node(step4) builder.add_edge(START, "planner") - graph = builder.compile(checkpointer=checkpointer) + graph = builder.compile(checkpointer=sync_checkpointer) config = {"configurable": {"thread_id": "1"}} @@ -5910,14 +5783,11 @@ def test_multistep_plan(request: pytest.FixtureRequest, checkpointer_name: str): } -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_command_goto_with_static_breakpoints( - request: pytest.FixtureRequest, checkpointer_name: str + sync_checkpointer: BaseCheckpointSaver, ) -> None: """Use Command goto with static breakpoints.""" - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - class State(TypedDict): """The graph state.""" @@ -5939,7 +5809,7 @@ def test_command_goto_with_static_breakpoints( builder.add_edge(START, "node1") builder.add_edge("node1", "node2") - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["node1"]) + graph = builder.compile(checkpointer=sync_checkpointer, interrupt_before=["node1"]) config = {"configurable": {"thread_id": str(uuid.uuid4())}} @@ -5982,14 +5852,11 @@ def test_parallel_node_execution(): assert duration < 3.0 -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_multiple_interrupt_state_persistence( - request: pytest.FixtureRequest, checkpointer_name: str + sync_checkpointer: BaseCheckpointSaver, ) -> None: """Test that state is preserved correctly across multiple interrupts.""" - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - class State(TypedDict): steps: Annotated[list[str], operator.add] @@ -6002,7 +5869,7 @@ def test_multiple_interrupt_state_persistence( builder.add_node("node", interruptible_node) builder.add_edge(START, "node") - app = builder.compile(checkpointer=checkpointer) + app = builder.compile(checkpointer=sync_checkpointer) config = {"configurable": {"thread_id": "1"}} # First execution - should hit first interrupt @@ -6171,18 +6038,16 @@ def test_multiple_updates() -> None: ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_falsy_return_from_task( - request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion + sync_checkpointer: BaseCheckpointSaver, snapshot: SnapshotAssertion ): """Test with a falsy return from a task.""" - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") @task def falsy_task() -> bool: return False - @entrypoint(checkpointer=checkpointer) + @entrypoint(checkpointer=sync_checkpointer) def graph(state: dict) -> dict: """React tool.""" falsy_task().result() @@ -6443,12 +6308,8 @@ def test_falsy_return_from_task( ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_multiple_interrupts_functional( - request: pytest.FixtureRequest, checkpointer_name: str -): +def test_multiple_interrupts_functional(sync_checkpointer: BaseCheckpointSaver): """Test multiple interrupts with functional API.""" - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") counter = 0 @@ -6459,7 +6320,7 @@ def test_multiple_interrupts_functional( counter += 1 return 2 * x - @entrypoint(checkpointer=checkpointer) + @entrypoint(checkpointer=sync_checkpointer) def graph(state: dict) -> dict: """React tool.""" @@ -6482,12 +6343,10 @@ def test_multiple_interrupts_functional( assert counter == 3 -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_multiple_interrupts_functional_cache( - request: pytest.FixtureRequest, checkpointer_name: str, cache: BaseCache + sync_checkpointer: BaseCheckpointSaver, cache: BaseCache ): """Test multiple interrupts with functional API.""" - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") counter = 0 @@ -6498,7 +6357,7 @@ def test_multiple_interrupts_functional_cache( counter += 1 return 2 * x - @entrypoint(checkpointer=checkpointer, cache=cache) + @entrypoint(checkpointer=sync_checkpointer, cache=cache) def graph(state: dict) -> dict: """React tool.""" @@ -6557,12 +6416,7 @@ def test_multiple_interrupts_functional_cache( assert counter == 6 -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_double_interrupt_subgraph( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - +def test_double_interrupt_subgraph(sync_checkpointer: BaseCheckpointSaver) -> None: class AgentState(TypedDict): input: str @@ -6584,7 +6438,7 @@ def test_double_interrupt_subgraph( ) # invoke the sub graph - subgraph = subgraph_builder.compile(checkpointer=checkpointer) + subgraph = subgraph_builder.compile(checkpointer=sync_checkpointer) thread = {"configurable": {"thread_id": str(uuid.uuid4())}} assert [c for c in subgraph.stream({"input": "test"}, thread)] == [ { @@ -6632,7 +6486,7 @@ def test_double_interrupt_subgraph( .add_node("invoke_sub_agent", invoke_sub_agent) .add_edge(START, "invoke_sub_agent") .add_edge("invoke_sub_agent", END) - .compile(checkpointer=checkpointer) + .compile(checkpointer=sync_checkpointer) ) assert [c for c in parent_agent.stream({"input": "test"}, thread)] == [ @@ -6670,10 +6524,7 @@ def test_double_interrupt_subgraph( ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_multi_resume(request: pytest.FixtureRequest, checkpointer_name: str) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - +def test_multi_resume(sync_checkpointer: BaseCheckpointSaver) -> None: class ChildState(TypedDict): prompt: str human_input: str @@ -6692,7 +6543,7 @@ def test_multi_resume(request: pytest.FixtureRequest, checkpointer_name: str) -> .add_node("get_human_input", get_human_input) .add_edge(START, "get_human_input") .add_edge("get_human_input", END) - .compile(checkpointer=checkpointer) + .compile(checkpointer=sync_checkpointer) ) class ParentState(TypedDict): @@ -6718,7 +6569,7 @@ def test_multi_resume(request: pytest.FixtureRequest, checkpointer_name: str) -> .add_conditional_edges(START, assign_workers, ["child_graph"]) .add_edge("child_graph", "cleanup") .add_edge("cleanup", END) - .compile(checkpointer=checkpointer) + .compile(checkpointer=sync_checkpointer) ) thread_config: RunnableConfig = { @@ -6912,12 +6763,7 @@ def test_entrypoint_from_sync_generator() -> None: yield "b" -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_multiple_subgraphs( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - +def test_multiple_subgraphs(sync_checkpointer: BaseCheckpointSaver) -> None: class State(TypedDict): a: int b: int @@ -6953,7 +6799,7 @@ def test_multiple_subgraphs( StateGraph(State, output=Output) .add_node(call_same_subgraph) .add_edge(START, "call_same_subgraph") - .compile(checkpointer=checkpointer) + .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} assert parent_call_same_subgraph.invoke({"a": 2, "b": 3}, config) == {"result": 15} @@ -6975,7 +6821,7 @@ def test_multiple_subgraphs( StateGraph(State, output=Output) .add_node(call_multiple_subgraphs) .add_edge(START, "call_multiple_subgraphs") - .compile(checkpointer=checkpointer) + .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "2"}} assert parent_call_multiple_subgraphs.invoke({"a": 2, "b": 3}, config) == { @@ -6984,12 +6830,7 @@ def test_multiple_subgraphs( } -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_multiple_subgraphs_functional( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - +def test_multiple_subgraphs_functional(sync_checkpointer: BaseCheckpointSaver) -> None: # Define addition subgraph @entrypoint() def add(inputs: tuple[int, int]): @@ -7012,7 +6853,7 @@ def test_multiple_subgraphs_functional( another_result = add.invoke([result, 10]) return another_result - @entrypoint(checkpointer=checkpointer) + @entrypoint(checkpointer=sync_checkpointer) def parent_call_same_subgraph(inputs): return call_same_subgraph(*inputs).result() @@ -7026,7 +6867,7 @@ def test_multiple_subgraphs_functional( multiply_result = multiply.invoke([a, b]) return [add_result, multiply_result] - @entrypoint(checkpointer=checkpointer) + @entrypoint(checkpointer=sync_checkpointer) def parent_call_multiple_subgraphs(inputs): return call_multiple_subgraphs(*inputs).result() @@ -7034,12 +6875,10 @@ def test_multiple_subgraphs_functional( assert parent_call_multiple_subgraphs.invoke([2, 3], config) == [5, 6] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_multiple_subgraphs_mixed_entrypoint( - request: pytest.FixtureRequest, checkpointer_name: str + sync_checkpointer: BaseCheckpointSaver, ) -> None: """Test calling multiple StateGraph subgraphs from an entrypoint.""" - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") class State(TypedDict): a: int @@ -7073,7 +6912,7 @@ def test_multiple_subgraphs_mixed_entrypoint( another_result = add_subgraph.invoke({"a": result, "b": 10})["result"] return another_result - @entrypoint(checkpointer=checkpointer) + @entrypoint(checkpointer=sync_checkpointer) def parent_call_same_subgraph(inputs): return call_same_subgraph(*inputs).result() @@ -7087,7 +6926,7 @@ def test_multiple_subgraphs_mixed_entrypoint( multiply_result = multiply_subgraph.invoke({"a": a, "b": b})["result"] return [add_result, multiply_result] - @entrypoint(checkpointer=checkpointer) + @entrypoint(checkpointer=sync_checkpointer) def parent_call_multiple_subgraphs(inputs): return call_multiple_subgraphs(*inputs).result() @@ -7095,12 +6934,10 @@ def test_multiple_subgraphs_mixed_entrypoint( assert parent_call_multiple_subgraphs.invoke([2, 3], config) == [5, 6] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_multiple_subgraphs_mixed_state_graph( - request: pytest.FixtureRequest, checkpointer_name: str + sync_checkpointer: BaseCheckpointSaver, ) -> None: """Test calling multiple entrypoint "subgraphs" from a StateGraph.""" - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") class State(TypedDict): a: int @@ -7134,7 +6971,7 @@ def test_multiple_subgraphs_mixed_state_graph( StateGraph(State, output=Output) .add_node(call_same_subgraph) .add_edge(START, "call_same_subgraph") - .compile(checkpointer=checkpointer) + .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} assert parent_call_same_subgraph.invoke({"a": 2, "b": 3}, config) == {"result": 15} @@ -7156,7 +6993,7 @@ def test_multiple_subgraphs_mixed_state_graph( StateGraph(State, output=Output) .add_node(call_multiple_subgraphs) .add_edge(START, "call_multiple_subgraphs") - .compile(checkpointer=checkpointer) + .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "2"}} assert parent_call_multiple_subgraphs.invoke({"a": 2, "b": 3}, config) == { @@ -7165,12 +7002,9 @@ def test_multiple_subgraphs_mixed_state_graph( } -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_multiple_subgraphs_checkpointer( - request: pytest.FixtureRequest, checkpointer_name: str + sync_checkpointer: BaseCheckpointSaver, ) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - class SubgraphState(TypedDict): sub_counter: Annotated[int, operator.add] @@ -7209,7 +7043,7 @@ def test_multiple_subgraphs_checkpointer( StateGraph(ParentState) .add_node(parent_node) .add_edge(START, "parent_node") - .compile(checkpointer=checkpointer) + .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} @@ -7708,13 +7542,9 @@ def test_stream_messages_dedupe_inputs() -> None: assert chunks[0][1]["langgraph_node"] == "call_model" -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_stream_messages_dedupe_state( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: +def test_stream_messages_dedupe_state(sync_checkpointer: BaseCheckpointSaver) -> None: from langchain_core.messages import AIMessage - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") to_emit = [AIMessage("bye", id="1"), AIMessage("bye again", id="2")] def call_model(state): @@ -7737,7 +7567,7 @@ def test_stream_messages_dedupe_state( .add_node("node_1", subgraph) .add_node("node_2", lambda state: state) .add_edge(START, "node_1") - .compile(checkpointer=checkpointer) + .compile(checkpointer=sync_checkpointer) ) thread1 = {"configurable": {"thread_id": "1"}} @@ -7768,12 +7598,9 @@ def test_stream_messages_dedupe_state( assert chunks[0][1]["langgraph_node"] == "call_model" -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_interrupt_subgraph_reenter_checkpointer_true( - request: pytest.FixtureRequest, checkpointer_name: str + sync_checkpointer: BaseCheckpointSaver, ) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - class SubgraphState(TypedDict): foo: str bar: str @@ -7824,7 +7651,7 @@ def test_interrupt_subgraph_reenter_checkpointer_true( .add_node(node) .add_edge(START, "call_subgraph") .add_edge("call_subgraph", "node") - .compile(checkpointer=checkpointer) + .compile(checkpointer=sync_checkpointer) ) config = {"configurable": {"thread_id": "1"}} @@ -7933,12 +7760,7 @@ def test_empty_invoke() -> None: } -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_parallel_interrupts( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - +def test_parallel_interrupts(sync_checkpointer: BaseCheckpointSaver) -> None: # --- CHILD GRAPH --- class ChildState(BaseModel): @@ -7994,7 +7816,7 @@ def test_parallel_interrupts( parent_graph_builder.add_edge("child_graph", "cleanup") parent_graph_builder.add_edge("cleanup", END) - parent_graph = parent_graph_builder.compile(checkpointer=checkpointer) + parent_graph = parent_graph_builder.compile(checkpointer=sync_checkpointer) # --- CLIENT INVOCATION --- @@ -8108,12 +7930,7 @@ def test_parallel_interrupts( ) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_parallel_interrupts_double( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - +def test_parallel_interrupts_double(sync_checkpointer: BaseCheckpointSaver) -> None: # --- CHILD GRAPH --- class ChildState(BaseModel): @@ -8176,7 +7993,7 @@ def test_parallel_interrupts_double( parent_graph_builder.add_edge("child_graph", "cleanup") parent_graph_builder.add_edge("cleanup", END) - parent_graph = parent_graph_builder.compile(checkpointer=checkpointer) + parent_graph = parent_graph_builder.compile(checkpointer=sync_checkpointer) # --- CLIENT INVOCATION --- diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index da2ab68dd..5756d6183 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -63,11 +63,6 @@ from langgraph.types import ( interrupt, ) from tests.any_str import AnyStr, AnyVersion, FloatBetween, UnsortedSequence -from tests.conftest import ( - ALL_CHECKPOINTERS_ASYNC, - SHOULD_CHECK_SNAPSHOTS, - awith_checkpointer, -) from tests.fake_tracer import FakeTracer from tests.memory_assert import MemorySaverNoPending from tests.messages import ( @@ -499,8 +494,7 @@ async def test_node_cancellation_on_other_node_exception_two() -> None: @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_dynamic_interrupt(checkpointer_name: str) -> None: +async def test_dynamic_interrupt(async_checkpointer: BaseCheckpointSaver) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -543,144 +537,136 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None: "market": "US", } - async with awith_checkpointer(checkpointer_name) as checkpointer: - tool_two = tool_two_graph.compile(checkpointer=checkpointer) + tool_two = tool_two_graph.compile(checkpointer=async_checkpointer) - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - # flow: interrupt -> resume with answer - thread2 = {"configurable": {"thread_id": "2"}} - # stop when about to enter node - assert [ - c - async for c in tool_two.astream( - {"my_key": "value ⛰️", "market": "DE"}, thread2 - ) - ] == [ - { - "__interrupt__": ( - Interrupt( - value="Just because...", - resumable=True, - ns=[AnyStr("tool_two:")], - ), - ) - }, - ] - # resume with answer - assert [ - c async for c in tool_two.astream(Command(resume=" my answer"), thread2) - ] == [ - {"tool_two": {"my_key": " my answer"}}, - ] - - # flow: interrupt -> clear - thread1 = {"configurable": {"thread_id": "1"}} - # stop when about to enter node - assert [ - c - async for c in tool_two.astream( - {"my_key": "value ⛰️", "market": "DE"}, thread1 - ) - ] == [ - { - "__interrupt__": ( - Interrupt( - value="Just because...", - resumable=True, - ns=[AnyStr("tool_two:")], - ), - ) - }, - ] - assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, - "thread_id": "1", - }, - ] - tup = await tool_two.checkpointer.aget_tuple(thread1) - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️", "market": "DE"}, - next=("tool_two",), - tasks=( - PregelTask( - AnyStr(), - "tool_two", - (PULL, "tool_two"), - interrupts=( - Interrupt( - value="Just because...", - resumable=True, - ns=[AnyStr("tool_two:")], - ), - ), - ), - ), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - }, - parent_config=( - [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config - ), - interrupts=( + # flow: interrupt -> resume with answer + thread2 = {"configurable": {"thread_id": "2"}} + # stop when about to enter node + assert [ + c + async for c in tool_two.astream({"my_key": "value ⛰️", "market": "DE"}, thread2) + ] == [ + { + "__interrupt__": ( Interrupt( value="Just because...", resumable=True, ns=[AnyStr("tool_two:")], ), - ), - ) + ) + }, + ] + # resume with answer + assert [ + c async for c in tool_two.astream(Command(resume=" my answer"), thread2) + ] == [ + {"tool_two": {"my_key": " my answer"}}, + ] - # clear the interrupt and next tasks - await tool_two.aupdate_state(thread1, None, as_node=END) - # interrupt is cleared, as well as the next tasks - tup = await tool_two.checkpointer.aget_tuple(thread1) - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️", "market": "DE"}, - next=(), - tasks=(), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": {}, - "thread_id": "1", - }, - parent_config=( - [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config + # flow: interrupt -> clear + thread1 = {"configurable": {"thread_id": "1"}} + # stop when about to enter node + assert [ + c + async for c in tool_two.astream({"my_key": "value ⛰️", "market": "DE"}, thread1) + ] == [ + { + "__interrupt__": ( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:")], + ), + ) + }, + ] + assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, + { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, + "thread_id": "1", + }, + ] + tup = await tool_two.checkpointer.aget_tuple(thread1) + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=("tool_two",), + tasks=( + PregelTask( + AnyStr(), + "tool_two", + (PULL, "tool_two"), + interrupts=( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:")], + ), + ), ), - interrupts=(), - ) + ), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, + parent_config=( + [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][-1].config + ), + interrupts=( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:")], + ), + ), + ) + + # clear the interrupt and next tasks + await tool_two.aupdate_state(thread1, None, as_node=END) + # interrupt is cleared, as well as the next tasks + tup = await tool_two.checkpointer.aget_tuple(thread1) + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=(), + tasks=(), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": {}, + "thread_id": "1", + }, + parent_config=( + [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][-1].config + ), + interrupts=(), + ) @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_dynamic_interrupt_subgraph(checkpointer_name: str) -> None: +async def test_dynamic_interrupt_subgraph( + async_checkpointer: BaseCheckpointSaver, +) -> None: class SubgraphState(TypedDict): my_key: str market: str @@ -735,151 +721,145 @@ async def test_dynamic_interrupt_subgraph(checkpointer_name: str) -> None: "market": "US", } - async with awith_checkpointer(checkpointer_name) as checkpointer: - tool_two = tool_two_graph.compile(checkpointer=checkpointer) + tool_two = tool_two_graph.compile(checkpointer=async_checkpointer) - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - # flow: interrupt -> resume with answer - thread2 = {"configurable": {"thread_id": "2"}} - # stop when about to enter node - assert [ - c - async for c in tool_two.astream( - {"my_key": "value ⛰️", "market": "DE"}, thread2 - ) - ] == [ - { - "__interrupt__": ( - Interrupt( - value="Just because...", - resumable=True, - ns=[AnyStr("tool_two:"), AnyStr("do:")], - ), - ) - }, - ] - # resume with answer - assert [ - c async for c in tool_two.astream(Command(resume=" my answer"), thread2) - ] == [ - {"tool_two": {"my_key": " my answer", "market": "DE"}}, - ] - - # flow: interrupt -> clear - thread1 = {"configurable": {"thread_id": "1"}} - thread1root = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}} - # stop when about to enter node - assert [ - c - async for c in tool_two.astream( - {"my_key": "value ⛰️", "market": "DE"}, thread1 - ) - ] == [ - { - "__interrupt__": ( - Interrupt( - value="Just because...", - resumable=True, - ns=[AnyStr("tool_two:"), AnyStr("do:")], - ), - ) - }, - ] - assert [c.metadata async for c in tool_two.checkpointer.alist(thread1root)] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, - "thread_id": "1", - }, - ] - tup = await tool_two.checkpointer.aget_tuple(thread1) - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️", "market": "DE"}, - next=("tool_two",), - tasks=( - PregelTask( - AnyStr(), - "tool_two", - (PULL, "tool_two"), - interrupts=( - Interrupt( - value="Just because...", - resumable=True, - ns=[AnyStr("tool_two:"), AnyStr("do:")], - ), - ), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("tool_two:"), - } - }, - ), - ), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - }, - parent_config=( - [c async for c in tool_two.checkpointer.alist(thread1root, limit=2)][ - -1 - ].config - ), - interrupts=( + # flow: interrupt -> resume with answer + thread2 = {"configurable": {"thread_id": "2"}} + # stop when about to enter node + assert [ + c + async for c in tool_two.astream({"my_key": "value ⛰️", "market": "DE"}, thread2) + ] == [ + { + "__interrupt__": ( Interrupt( value="Just because...", resumable=True, ns=[AnyStr("tool_two:"), AnyStr("do:")], ), - ), - ) + ) + }, + ] + # resume with answer + assert [ + c async for c in tool_two.astream(Command(resume=" my answer"), thread2) + ] == [ + {"tool_two": {"my_key": " my answer", "market": "DE"}}, + ] - # clear the interrupt and next tasks - await tool_two.aupdate_state(thread1, None, as_node=END) - # interrupt is cleared, as well as the next tasks - tup = await tool_two.checkpointer.aget_tuple(thread1) - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️", "market": "DE"}, - next=(), - tasks=(), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "update", - "step": 1, - "writes": {}, - "thread_id": "1", - }, - parent_config=( - [c async for c in tool_two.checkpointer.alist(thread1root, limit=2)][ - -1 - ].config + # flow: interrupt -> clear + thread1 = {"configurable": {"thread_id": "1"}} + thread1root = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}} + # stop when about to enter node + assert [ + c + async for c in tool_two.astream({"my_key": "value ⛰️", "market": "DE"}, thread1) + ] == [ + { + "__interrupt__": ( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:"), AnyStr("do:")], + ), + ) + }, + ] + assert [c.metadata async for c in tool_two.checkpointer.alist(thread1root)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, + { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, + "thread_id": "1", + }, + ] + tup = await tool_two.checkpointer.aget_tuple(thread1) + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=("tool_two",), + tasks=( + PregelTask( + AnyStr(), + "tool_two", + (PULL, "tool_two"), + interrupts=( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:"), AnyStr("do:")], + ), + ), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("tool_two:"), + } + }, ), - interrupts=(), - ) + ), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, + parent_config=( + [c async for c in tool_two.checkpointer.alist(thread1root, limit=2)][ + -1 + ].config + ), + interrupts=( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:"), AnyStr("do:")], + ), + ), + ) + + # clear the interrupt and next tasks + await tool_two.aupdate_state(thread1, None, as_node=END) + # interrupt is cleared, as well as the next tasks + tup = await tool_two.checkpointer.aget_tuple(thread1) + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=(), + tasks=(), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": {}, + "thread_id": "1", + }, + parent_config=( + [c async for c in tool_two.checkpointer.alist(thread1root, limit=2)][ + -1 + ].config + ), + interrupts=(), + ) @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_copy_checkpoint(checkpointer_name: str) -> None: +async def test_copy_checkpoint(async_checkpointer: BaseCheckpointSaver) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] market: str @@ -929,173 +909,165 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None: "market": "US", } - async with awith_checkpointer(checkpointer_name) as checkpointer: - tool_two = tool_two_graph.compile(checkpointer=checkpointer) + tool_two = tool_two_graph.compile(checkpointer=async_checkpointer) - # missing thread_id - with pytest.raises(ValueError, match="thread_id"): - await tool_two.ainvoke({"my_key": "value", "market": "DE"}) + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + await tool_two.ainvoke({"my_key": "value", "market": "DE"}) - # flow: interrupt -> resume with answer - thread2 = {"configurable": {"thread_id": "2"}} - # stop when about to enter node - assert [ - c - async for c in tool_two.astream( - {"my_key": "value ⛰️", "market": "DE"}, thread2 + # flow: interrupt -> resume with answer + thread2 = {"configurable": {"thread_id": "2"}} + # stop when about to enter node + assert [ + c + async for c in tool_two.astream({"my_key": "value ⛰️", "market": "DE"}, thread2) + ] == UnsortedSequence( + { + "__interrupt__": ( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:")], + ), ) - ] == UnsortedSequence( - { - "__interrupt__": ( + }, + { + "tool_one": {"my_key": " one"}, + }, + ) + # resume with answer + assert [ + c async for c in tool_two.astream(Command(resume=" my answer"), thread2) + ] == [ + { + "__metadata__": {"cached": True}, + "tool_one": {"my_key": " one"}, + }, + {"tool_two": {"my_key": " my answer"}}, + ] + + # flow: interrupt -> clear tasks + thread1 = {"configurable": {"thread_id": "1"}} + # stop when about to enter node + assert await tool_two.ainvoke({"my_key": "value ⛰️", "market": "DE"}, thread1) == { + "my_key": "value ⛰️ one", + "market": "DE", + "__interrupt__": [ + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:")], + ) + ], + } + + assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, + { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, + "thread_id": "1", + }, + ] + + tup = await tool_two.checkpointer.aget_tuple(thread1) + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️ one", "market": "DE"}, + next=("tool_two",), + tasks=( + PregelTask( + AnyStr(), + name="tool_one", + path=("__pregel_push", 0, False), + error=None, + interrupts=(), + state=None, + result={"my_key": " one"}, + ), + PregelTask( + AnyStr(), + "tool_two", + (PULL, "tool_two"), + interrupts=( Interrupt( value="Just because...", resumable=True, ns=[AnyStr("tool_two:")], ), - ) - }, - { - "tool_one": {"my_key": " one"}, - }, - ) - # resume with answer - assert [ - c async for c in tool_two.astream(Command(resume=" my answer"), thread2) - ] == [ - { - "__metadata__": {"cached": True}, - "tool_one": {"my_key": " one"}, - }, - {"tool_two": {"my_key": " my answer"}}, - ] + ), + ), + ), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, + parent_config=( + [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][-1].config + ), + interrupts=( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:")], + ), + ), + ) - # flow: interrupt -> clear tasks - thread1 = {"configurable": {"thread_id": "1"}} - # stop when about to enter node - assert await tool_two.ainvoke( - {"my_key": "value ⛰️", "market": "DE"}, thread1 - ) == { - "my_key": "value ⛰️ one", - "market": "DE", - "__interrupt__": [ - Interrupt( - value="Just because...", - resumable=True, - ns=[AnyStr("tool_two:")], - ) - ], - } - - assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ - { - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - }, - { - "parents": {}, - "source": "input", - "step": -1, - "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, - "thread_id": "1", - }, - ] - - tup = await tool_two.checkpointer.aget_tuple(thread1) - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️ one", "market": "DE"}, - next=("tool_two",), - tasks=( - PregelTask( - AnyStr(), - name="tool_one", - path=("__pregel_push", 0, False), - error=None, - interrupts=(), - state=None, - result={"my_key": " one"}, - ), - PregelTask( - AnyStr(), - "tool_two", - (PULL, "tool_two"), - interrupts=( - Interrupt( - value="Just because...", - resumable=True, - ns=[AnyStr("tool_two:")], - ), - ), - ), + # clear the interrupt and next tasks + await tool_two.aupdate_state(thread1, None, as_node="__copy__") + # interrupt is cleared, next task is kept + tup = await tool_two.checkpointer.aget_tuple(thread1) + assert await tool_two.aget_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=("tool_one", "tool_two"), + tasks=( + PregelTask( + AnyStr(), + "tool_one", + (PUSH, 0, False), + result=None, ), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, - "thread_id": "1", - }, - parent_config=( - [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config + PregelTask( + AnyStr(), + "tool_two", + (PULL, "tool_two"), + interrupts=(), ), - interrupts=( - Interrupt( - value="Just because...", - resumable=True, - ns=[AnyStr("tool_two:")], - ), - ), - ) - - # clear the interrupt and next tasks - await tool_two.aupdate_state(thread1, None, as_node="__copy__") - # interrupt is cleared, next task is kept - tup = await tool_two.checkpointer.aget_tuple(thread1) - assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️", "market": "DE"}, - next=("tool_one", "tool_two"), - tasks=( - PregelTask( - AnyStr(), - "tool_one", - (PUSH, 0, False), - result=None, - ), - PregelTask( - AnyStr(), - "tool_two", - (PULL, "tool_two"), - interrupts=(), - ), - ), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "fork", - "step": 1, - "writes": None, - "thread_id": "1", - }, - parent_config=( - [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].parent_config - ), - interrupts=(), - ) + ), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "fork", + "step": 1, + "writes": None, + "thread_id": "1", + }, + parent_config=( + [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ + -1 + ].parent_config + ), + interrupts=(), + ) @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_node_not_cancelled_on_other_node_interrupted( - checkpointer_name: str, + async_checkpointer: BaseCheckpointSaver, ) -> None: class State(TypedDict): hello: Annotated[str, operator.add] @@ -1123,46 +1095,45 @@ async def test_node_not_cancelled_on_other_node_interrupted( builder.add_node("bad", iambad) builder.set_conditional_entry_point(lambda _: ["agent", "bad"], then=END) - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) - thread = {"configurable": {"thread_id": "1"}} + graph = builder.compile(checkpointer=async_checkpointer) + thread = {"configurable": {"thread_id": "1"}} - # writes from "awhile" are applied to last chunk - assert await graph.ainvoke({"hello": "world"}, thread) == { - "hello": "world again", - "__interrupt__": [ - Interrupt( - value="I am bad", - resumable=True, - ns=[AnyStr("bad:")], - ) - ], - } + # writes from "awhile" are applied to last chunk + assert await graph.ainvoke({"hello": "world"}, thread) == { + "hello": "world again", + "__interrupt__": [ + Interrupt( + value="I am bad", + resumable=True, + ns=[AnyStr("bad:")], + ) + ], + } - assert not inner_task_cancelled - assert awhiles == 1 + assert not inner_task_cancelled + assert awhiles == 1 - assert await graph.ainvoke(None, thread, debug=True) == { - "hello": "world again", - "__interrupt__": [ - Interrupt( - value="I am bad", - resumable=True, - ns=[AnyStr("bad:")], - ) - ], - } + assert await graph.ainvoke(None, thread, debug=True) == { + "hello": "world again", + "__interrupt__": [ + Interrupt( + value="I am bad", + resumable=True, + ns=[AnyStr("bad:")], + ) + ], + } - assert not inner_task_cancelled - assert awhiles == 1 + assert not inner_task_cancelled + assert awhiles == 1 - # resume with answer - assert await graph.ainvoke(Command(resume=" okay"), thread) == { - "hello": "world again okay" - } + # resume with answer + assert await graph.ainvoke(Command(resume=" okay"), thread) == { + "hello": "world again okay" + } - assert not inner_task_cancelled - assert awhiles == 1 + assert not inner_task_cancelled + assert awhiles == 1 @pytest.mark.parametrize("stream_hang_s", [0.3, 0.6]) @@ -1456,15 +1427,14 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: graph.set_finish_point("add_one") gapp = graph.compile() - if SHOULD_CHECK_SNAPSHOTS: - assert app.input_schema.model_json_schema() == { - "title": "LangGraphInput", - "type": "integer", - } - assert app.output_schema.model_json_schema() == { - "title": "LangGraphOutput", - "type": "integer", - } + assert app.input_schema.model_json_schema() == { + "title": "LangGraphInput", + "type": "integer", + } + assert app.output_schema.model_json_schema() == { + "title": "LangGraphOutput", + "type": "integer", + } assert await app.ainvoke(2) == 3 assert await app.ainvoke(2, output_keys=["output"]) == {"output": 3} @@ -1504,24 +1474,23 @@ async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> N input_channels="input", ) - if SHOULD_CHECK_SNAPSHOTS: - assert app.input_schema.model_json_schema() == { - "title": "LangGraphInput", - "type": "integer", - } - assert app.output_schema.model_json_schema() == { - "title": "LangGraphOutput", - "type": "object", - "properties": { - "output": {"title": "Output", "type": "integer", "default": None}, - "fixed": {"title": "Fixed", "type": "integer", "default": None}, - "output_plus_one": { - "title": "Output Plus One", - "type": "integer", - "default": None, - }, + assert app.input_schema.model_json_schema() == { + "title": "LangGraphInput", + "type": "integer", + } + assert app.output_schema.model_json_schema() == { + "title": "LangGraphOutput", + "type": "object", + "properties": { + "output": {"title": "Output", "type": "integer", "default": None}, + "fixed": {"title": "Fixed", "type": "integer", "default": None}, + "output_plus_one": { + "title": "Output Plus One", + "type": "integer", + "default": None, }, - } + }, + } assert await app.ainvoke(2) == {"output": 3, "fixed": 5, "output_plus_one": 4} @@ -1536,18 +1505,17 @@ async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: output_channels=["output"], ) - if SHOULD_CHECK_SNAPSHOTS: - assert app.input_schema.model_json_schema() == { - "title": "LangGraphInput", - "type": "integer", - } - assert app.output_schema.model_json_schema() == { - "title": "LangGraphOutput", - "type": "object", - "properties": { - "output": {"title": "Output", "type": "integer", "default": None} - }, - } + assert app.input_schema.model_json_schema() == { + "title": "LangGraphInput", + "type": "integer", + } + assert app.output_schema.model_json_schema() == { + "title": "LangGraphOutput", + "type": "object", + "properties": { + "output": {"title": "Output", "type": "integer", "default": None} + }, + } assert await app.ainvoke(2) == {"output": 3} @@ -1562,21 +1530,18 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> output_channels=["output"], ) - if SHOULD_CHECK_SNAPSHOTS: - assert app.input_schema.model_json_schema() == { - "title": "LangGraphInput", - "type": "object", - "properties": { - "input": {"title": "Input", "type": "integer", "default": None} - }, - } - assert app.output_schema.model_json_schema() == { - "title": "LangGraphOutput", - "type": "object", - "properties": { - "output": {"title": "Output", "type": "integer", "default": None} - }, - } + assert app.input_schema.model_json_schema() == { + "title": "LangGraphInput", + "type": "object", + "properties": {"input": {"title": "Input", "type": "integer", "default": None}}, + } + assert app.output_schema.model_json_schema() == { + "title": "LangGraphOutput", + "type": "object", + "properties": { + "output": {"title": "Output", "type": "integer", "default": None} + }, + } assert await app.ainvoke({"input": 2}) == {"output": 3} @@ -1909,8 +1874,9 @@ async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) assert await app.ainvoke(2) == [3, 3] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_invoke_checkpoint(mocker: MockerFixture, checkpointer_name: str) -> None: +async def test_invoke_checkpoint( + mocker: MockerFixture, async_checkpointer: BaseCheckpointSaver +) -> None: add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) errored_once = False @@ -1933,51 +1899,49 @@ async def test_invoke_checkpoint(mocker: MockerFixture, checkpointer_name: str) | raise_if_above_10 ) - async with awith_checkpointer(checkpointer_name) as checkpointer: - app = Pregel( - nodes={"one": one}, - channels={ - "total": BinaryOperatorAggregate(int, operator.add), - "input": LastValue(int), - "output": LastValue(int), - }, - input_channels="input", - output_channels="output", - checkpointer=checkpointer, - retry_policy=RetryPolicy(), - ) + app = Pregel( + nodes={"one": one}, + channels={ + "total": BinaryOperatorAggregate(int, operator.add), + "input": LastValue(int), + "output": LastValue(int), + }, + input_channels="input", + output_channels="output", + checkpointer=async_checkpointer, + retry_policy=RetryPolicy(), + ) - # total starts out as 0, so output is 0+2=2 - assert await app.ainvoke(2, {"configurable": {"thread_id": "1"}}) == 2 - checkpoint = await checkpointer.aget({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 2 - # total is now 2, so output is 2+3=5 - assert await app.ainvoke(3, {"configurable": {"thread_id": "1"}}) == 5 - assert errored_once, "errored and retried" - checkpoint = await checkpointer.aget({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 7 - # total is now 2+5=7, so output would be 7+4=11, but raises ValueError - with pytest.raises(ValueError): - await app.ainvoke(4, {"configurable": {"thread_id": "1"}}) - # checkpoint is not updated - checkpoint = await checkpointer.aget({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 7 - # on a new thread, total starts out as 0, so output is 0+5=5 - assert await app.ainvoke(5, {"configurable": {"thread_id": "2"}}) == 5 - checkpoint = await checkpointer.aget({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 7 - checkpoint = await checkpointer.aget({"configurable": {"thread_id": "2"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 5 + # total starts out as 0, so output is 0+2=2 + assert await app.ainvoke(2, {"configurable": {"thread_id": "1"}}) == 2 + checkpoint = await async_checkpointer.aget({"configurable": {"thread_id": "1"}}) + assert checkpoint is not None + assert checkpoint["channel_values"].get("total") == 2 + # total is now 2, so output is 2+3=5 + assert await app.ainvoke(3, {"configurable": {"thread_id": "1"}}) == 5 + assert errored_once, "errored and retried" + checkpoint = await async_checkpointer.aget({"configurable": {"thread_id": "1"}}) + assert checkpoint is not None + assert checkpoint["channel_values"].get("total") == 7 + # total is now 2+5=7, so output would be 7+4=11, but raises ValueError + with pytest.raises(ValueError): + await app.ainvoke(4, {"configurable": {"thread_id": "1"}}) + # checkpoint is not updated + checkpoint = await async_checkpointer.aget({"configurable": {"thread_id": "1"}}) + assert checkpoint is not None + assert checkpoint["channel_values"].get("total") == 7 + # on a new thread, total starts out as 0, so output is 0+5=5 + assert await app.ainvoke(5, {"configurable": {"thread_id": "2"}}) == 5 + checkpoint = await async_checkpointer.aget({"configurable": {"thread_id": "1"}}) + assert checkpoint is not None + assert checkpoint["channel_values"].get("total") == 7 + checkpoint = await async_checkpointer.aget({"configurable": {"thread_id": "2"}}) + assert checkpoint is not None + assert checkpoint["channel_values"].get("total") == 5 -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_pending_writes_resume( - checkpointer_name: str, checkpoint_during: bool + async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool ) -> None: class State(TypedDict): value: Annotated[int, operator.add] @@ -2008,244 +1972,237 @@ async def test_pending_writes_resume( ) builder.add_edge(START, "one") builder.add_edge(START, "two") - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) + graph = builder.compile(checkpointer=async_checkpointer) - thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} - with pytest.raises(ConnectionError, match="I'm not good"): - await graph.ainvoke( - {"value": 1}, thread1, checkpoint_during=checkpoint_during - ) + thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} + with pytest.raises(ConnectionError, match="I'm not good"): + await graph.ainvoke({"value": 1}, thread1, checkpoint_during=checkpoint_during) - # both nodes should have been called once - assert one.calls == 1 - assert two.calls == 2 + # both nodes should have been called once + assert one.calls == 1 + assert two.calls == 2 - # latest checkpoint should be before nodes "one", "two" - # but we should have applied pending writes from "one" - state = await graph.aget_state(thread1) - assert state is not None - assert state.values == {"value": 3} - assert state.next == ("two",) - assert state.tasks == ( - PregelTask(AnyStr(), "one", (PULL, "one"), result={"value": 2}), - PregelTask( - AnyStr(), - "two", - (PULL, "two"), - 'ConnectionError("I\'m not good")', - ), - ) - assert state.metadata == { + # latest checkpoint should be before nodes "one", "two" + # but we should have applied pending writes from "one" + state = await graph.aget_state(thread1) + assert state is not None + assert state.values == {"value": 3} + assert state.next == ("two",) + assert state.tasks == ( + PregelTask(AnyStr(), "one", (PULL, "one"), result={"value": 2}), + PregelTask( + AnyStr(), + "two", + (PULL, "two"), + 'ConnectionError("I\'m not good")', + ), + ) + assert state.metadata == { + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + } + # get_state with checkpoint_id should not apply any pending writes + state = await graph.aget_state(state.config) + assert state is not None + assert state.values == {"value": 1} + assert state.next == ("one", "two") + # should contain pending write of "one" + checkpoint = await async_checkpointer.aget_tuple(thread1) + assert checkpoint is not None + # should contain error from "two" + expected_writes = [ + (AnyStr(), "value", 2), + (AnyStr(), ERROR, 'ConnectionError("I\'m not good")'), + ] + assert len(checkpoint.pending_writes) == 2 + assert all(w in expected_writes for w in checkpoint.pending_writes) + # both non-error pending writes come from same task + non_error_writes = [w for w in checkpoint.pending_writes if w[1] != ERROR] + # error write is from the other task + error_write = next(w for w in checkpoint.pending_writes if w[1] == ERROR) + assert error_write[0] != non_error_writes[0][0] + + # resume execution + with pytest.raises(ConnectionError, match="I'm not good"): + await graph.ainvoke(None, thread1, checkpoint_during=checkpoint_during) + + # node "one" succeeded previously, so shouldn't be called again + assert one.calls == 1 + # node "two" should have been called once again + assert two.calls == 4 + + # confirm no new checkpoints saved + state_two = await graph.aget_state(thread1) + assert state_two.metadata == state.metadata + + # resume execution, without exception + two.rtn = {"value": 3} + # both the pending write and the new write were applied, 1 + 2 + 3 = 6 + assert await graph.ainvoke(None, thread1, checkpoint_during=checkpoint_during) == { + "value": 6 + } + + # check all final checkpoints + checkpoints = [c async for c in async_checkpointer.alist(thread1)] + # we should have 3 + assert len(checkpoints) == (3 if checkpoint_during else 2) + # the last one not too interesting for this test + assert checkpoints[0] == CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + checkpoint={ + "v": 3, + "id": AnyStr(), + "ts": AnyStr(), + "pending_sends": [], + "versions_seen": { + "one": { + "branch:to:one": AnyVersion(), + }, + "two": { + "branch:to:two": AnyVersion(), + }, + "__input__": {}, + "__start__": { + "__start__": AnyVersion(), + }, + "__interrupt__": { + "value": AnyVersion(), + "__start__": AnyVersion(), + "branch:to:one": AnyVersion(), + "branch:to:two": AnyVersion(), + }, + }, + "channel_versions": { + "value": AnyVersion(), + "__start__": AnyVersion(), + "branch:to:one": AnyVersion(), + "branch:to:two": AnyVersion(), + }, + "channel_values": {"value": 6}, + }, + metadata={ "parents": {}, + "step": 1, "source": "loop", + "writes": {"one": {"value": 2}, "two": {"value": 3}}, + "thread_id": "1", + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": checkpoints[1].config["configurable"]["checkpoint_id"], + } + }, + pending_writes=[], + ) + # the previous one we assert that pending writes contains both + # - original error + # - successful writes from resuming after preventing error + assert checkpoints[1] == CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + checkpoint={ + "v": 3, + "id": AnyStr(), + "ts": AnyStr(), + "pending_sends": [], + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": AnyVersion(), + }, + }, + "channel_versions": { + "value": AnyVersion(), + "__start__": AnyVersion(), + "branch:to:one": AnyVersion(), + "branch:to:two": AnyVersion(), + }, + "channel_values": { + "value": 1, + "branch:to:one": None, + "branch:to:two": None, + }, + }, + metadata={ + "parents": {}, "step": 0, + "source": "loop", "writes": None, "thread_id": "1", - } - # get_state with checkpoint_id should not apply any pending writes - state = await graph.aget_state(state.config) - assert state is not None - assert state.values == {"value": 1} - assert state.next == ("one", "two") - # should contain pending write of "one" - checkpoint = await checkpointer.aget_tuple(thread1) - assert checkpoint is not None - # should contain error from "two" - expected_writes = [ + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"] + if checkpoint_during + else AnyStr(), + } + }, + pending_writes=UnsortedSequence( (AnyStr(), "value", 2), - (AnyStr(), ERROR, 'ConnectionError("I\'m not good")'), - ] - assert len(checkpoint.pending_writes) == 2 - assert all(w in expected_writes for w in checkpoint.pending_writes) - # both non-error pending writes come from same task - non_error_writes = [w for w in checkpoint.pending_writes if w[1] != ERROR] - # error write is from the other task - error_write = next(w for w in checkpoint.pending_writes if w[1] == ERROR) - assert error_write[0] != non_error_writes[0][0] - - # resume execution - with pytest.raises(ConnectionError, match="I'm not good"): - await graph.ainvoke(None, thread1, checkpoint_during=checkpoint_during) - - # node "one" succeeded previously, so shouldn't be called again - assert one.calls == 1 - # node "two" should have been called once again - assert two.calls == 4 - - # confirm no new checkpoints saved - state_two = await graph.aget_state(thread1) - assert state_two.metadata == state.metadata - - # resume execution, without exception - two.rtn = {"value": 3} - # both the pending write and the new write were applied, 1 + 2 + 3 = 6 - assert await graph.ainvoke( - None, thread1, checkpoint_during=checkpoint_during - ) == {"value": 6} - - # check all final checkpoints - checkpoints = [c async for c in checkpointer.alist(thread1)] - # we should have 3 - assert len(checkpoints) == (3 if checkpoint_during else 2) - # the last one not too interesting for this test - assert checkpoints[0] == CheckpointTuple( - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - checkpoint={ - "v": 3, - "id": AnyStr(), - "ts": AnyStr(), - "pending_sends": [], - "versions_seen": { - "one": { - "branch:to:one": AnyVersion(), - }, - "two": { - "branch:to:two": AnyVersion(), - }, - "__input__": {}, - "__start__": { - "__start__": AnyVersion(), - }, - "__interrupt__": { - "value": AnyVersion(), - "__start__": AnyVersion(), - "branch:to:one": AnyVersion(), - "branch:to:two": AnyVersion(), - }, - }, - "channel_versions": { - "value": AnyVersion(), - "__start__": AnyVersion(), - "branch:to:one": AnyVersion(), - "branch:to:two": AnyVersion(), - }, - "channel_values": {"value": 6}, - }, - metadata={ - "parents": {}, - "step": 1, - "source": "loop", - "writes": {"one": {"value": 2}, "two": {"value": 3}}, - "thread_id": "1", - }, - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": checkpoints[1].config["configurable"][ - "checkpoint_id" - ], - } - }, - pending_writes=[], + (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'), + (AnyStr(), "value", 3), ) - # the previous one we assert that pending writes contains both - # - original error - # - successful writes from resuming after preventing error - assert checkpoints[1] == CheckpointTuple( - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - checkpoint={ - "v": 3, - "id": AnyStr(), - "ts": AnyStr(), - "pending_sends": [], - "versions_seen": { - "__input__": {}, - "__start__": { - "__start__": AnyVersion(), - }, - }, - "channel_versions": { - "value": AnyVersion(), - "__start__": AnyVersion(), - "branch:to:one": AnyVersion(), - "branch:to:two": AnyVersion(), - }, - "channel_values": { - "value": 1, - "branch:to:one": None, - "branch:to:two": None, - }, - }, - metadata={ - "parents": {}, - "step": 0, - "source": "loop", - "writes": None, + if checkpoint_during + else UnsortedSequence( + (AnyStr(), "value", 2), + (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'), + # the write against the previous checkpoint is not saved, as it is + # produced in a run where only the next checkpoint (the last) is saved + ), + ) + if not checkpoint_during: + return + assert checkpoints[2] == CheckpointTuple( + config={ + "configurable": { "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + checkpoint={ + "v": 3, + "id": AnyStr(), + "ts": AnyStr(), + "pending_sends": [], + "versions_seen": {"__input__": {}}, + "channel_versions": { + "__start__": AnyVersion(), }, - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": checkpoints[2].config["configurable"][ - "checkpoint_id" - ] - if checkpoint_during - else AnyStr(), - } - }, - pending_writes=UnsortedSequence( - (AnyStr(), "value", 2), - (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'), - (AnyStr(), "value", 3), - ) - if checkpoint_during - else UnsortedSequence( - (AnyStr(), "value", 2), - (AnyStr(), "__error__", 'ConnectionError("I\'m not good")'), - # the write against the previous checkpoint is not saved, as it is - # produced in a run where only the next checkpoint (the last) is saved - ), - ) - if not checkpoint_during: - return - assert checkpoints[2] == CheckpointTuple( - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - checkpoint={ - "v": 3, - "id": AnyStr(), - "ts": AnyStr(), - "pending_sends": [], - "versions_seen": {"__input__": {}}, - "channel_versions": { - "__start__": AnyVersion(), - }, - "channel_values": {"__start__": {"value": 1}}, - }, - metadata={ - "parents": {}, - "step": -1, - "source": "input", - "writes": {"__start__": {"value": 1}}, - "thread_id": "1", - }, - parent_config=None, - pending_writes=UnsortedSequence( - (AnyStr(), "value", 1), - (AnyStr(), "branch:to:one", None), - (AnyStr(), "branch:to:two", None), - ), - ) + "channel_values": {"__start__": {"value": 1}}, + }, + metadata={ + "parents": {}, + "step": -1, + "source": "input", + "writes": {"__start__": {"value": 1}}, + "thread_id": "1", + }, + parent_config=None, + pending_writes=UnsortedSequence( + (AnyStr(), "value", 1), + (AnyStr(), "branch:to:one", None), + (AnyStr(), "branch:to:two", None), + ), + ) async def test_run_from_checkpoint_id_retains_previous_writes( @@ -2401,8 +2358,7 @@ async def test_concurrent_emit_sends() -> None: ) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_send_sequences(checkpointer_name: str) -> None: +async def test_send_sequences(async_checkpointer: BaseCheckpointSaver) -> None: class Node: def __init__(self, name: str): self.name = name @@ -2450,89 +2406,89 @@ async def test_send_sequences(checkpointer_name: str) -> None: "3", ] - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["3.1"]) - thread1 = {"configurable": {"thread_id": "1"}} - assert await graph.ainvoke(["0"], thread1) == [ - "0", - "1", - ] - assert await graph.ainvoke(None, thread1) == [ - "0", - "1", - "3.1", - "2|Command(goto=Send(node='2', arg=3))", - "2|Command(goto=Send(node='2', arg=4))", - "3", - "2|3", - "2|4", - "3", - ] + graph = builder.compile(checkpointer=async_checkpointer, interrupt_before=["3.1"]) + thread1 = {"configurable": {"thread_id": "1"}} + assert await graph.ainvoke(["0"], thread1) == [ + "0", + "1", + ] + assert await graph.ainvoke(None, thread1) == [ + "0", + "1", + "3.1", + "2|Command(goto=Send(node='2', arg=3))", + "2|Command(goto=Send(node='2', arg=4))", + "3", + "2|3", + "2|4", + "3", + ] @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_imp_task(checkpointer_name: str, checkpoint_during: bool) -> None: - async with awith_checkpointer(checkpointer_name) as checkpointer: - mapper_calls = 0 +async def test_imp_task( + async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool +) -> None: + mapper_calls = 0 - @task() - async def mapper(input: int) -> str: - nonlocal mapper_calls - mapper_calls += 1 - await asyncio.sleep(0.1 * input) - return str(input) * 2 + @task() + async def mapper(input: int) -> str: + nonlocal mapper_calls + mapper_calls += 1 + await asyncio.sleep(0.1 * input) + return str(input) * 2 - @entrypoint(checkpointer=checkpointer) - async def graph(input: list[int]) -> list[str]: - futures = [mapper(i) for i in input] - mapped = await asyncio.gather(*futures) - answer = interrupt("question") - return [m + answer for m in mapped] + @entrypoint(checkpointer=async_checkpointer) + async def graph(input: list[int]) -> list[str]: + futures = [mapper(i) for i in input] + mapped = await asyncio.gather(*futures) + answer = interrupt("question") + return [m + answer for m in mapped] - tracer = FakeTracer() - thread1 = {"configurable": {"thread_id": "1"}, "callbacks": [tracer]} - assert [ - c - async for c in graph.astream( - [0, 1], thread1, checkpoint_during=checkpoint_during + tracer = FakeTracer() + thread1 = {"configurable": {"thread_id": "1"}, "callbacks": [tracer]} + assert [ + c + async for c in graph.astream( + [0, 1], thread1, checkpoint_during=checkpoint_during + ) + ] == [ + {"mapper": "00"}, + {"mapper": "11"}, + { + "__interrupt__": ( + Interrupt( + value="question", + resumable=True, + ns=[AnyStr("graph:")], + when="during", + ), ) - ] == [ - {"mapper": "00"}, - {"mapper": "11"}, - { - "__interrupt__": ( - Interrupt( - value="question", - resumable=True, - ns=[AnyStr("graph:")], - when="during", - ), - ) - }, - ] - assert mapper_calls == 2 - assert len(tracer.runs) == 1 - assert len(tracer.runs[0].child_runs) == 1 - entrypoint_run = tracer.runs[0].child_runs[0] - assert entrypoint_run.name == "graph" - mapper_runs = [r for r in entrypoint_run.child_runs if r.name == "mapper"] - assert len(mapper_runs) == 2 - assert any(r.inputs == {"input": 0} for r in mapper_runs) - assert any(r.inputs == {"input": 1} for r in mapper_runs) + }, + ] + assert mapper_calls == 2 + assert len(tracer.runs) == 1 + assert len(tracer.runs[0].child_runs) == 1 + entrypoint_run = tracer.runs[0].child_runs[0] + assert entrypoint_run.name == "graph" + mapper_runs = [r for r in entrypoint_run.child_runs if r.name == "mapper"] + assert len(mapper_runs) == 2 + assert any(r.inputs == {"input": 0} for r in mapper_runs) + assert any(r.inputs == {"input": 1} for r in mapper_runs) - assert await graph.ainvoke( - Command(resume="answer"), thread1, checkpoint_during=checkpoint_during - ) == [ - "00answer", - "11answer", - ] - assert mapper_calls == 2 + assert await graph.ainvoke( + Command(resume="answer"), thread1, checkpoint_during=checkpoint_during + ) == [ + "00answer", + "11answer", + ] + assert mapper_calls == 2 @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_imp_nested(checkpointer_name: str, checkpoint_during: bool) -> None: +async def test_imp_nested( + async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool +) -> None: async def mynode(input: list[str]) -> list[str]: return [it + "a" for it in input] @@ -2550,195 +2506,187 @@ async def test_imp_nested(checkpointer_name: str, checkpoint_during: bool) -> No await asyncio.sleep(input / 100) return await submapper(input) * 2 - async with awith_checkpointer(checkpointer_name) as checkpointer: + @entrypoint(checkpointer=async_checkpointer) + async def graph(input: list[int]) -> list[str]: + futures = [mapper(i) for i in input] + mapped = await asyncio.gather(*futures) + answer = interrupt("question") + final = [m + answer for m in mapped] + return await add_a.ainvoke(final) - @entrypoint(checkpointer=checkpointer) - async def graph(input: list[int]) -> list[str]: - futures = [mapper(i) for i in input] - mapped = await asyncio.gather(*futures) - answer = interrupt("question") - final = [m + answer for m in mapped] - return await add_a.ainvoke(final) + assert graph.get_input_jsonschema() == { + "type": "array", + "items": {"type": "integer"}, + "title": "LangGraphInput", + } + assert graph.get_output_jsonschema() == { + "type": "array", + "items": {"type": "string"}, + "title": "LangGraphOutput", + } - assert graph.get_input_jsonschema() == { - "type": "array", - "items": {"type": "integer"}, - "title": "LangGraphInput", - } - assert graph.get_output_jsonschema() == { - "type": "array", - "items": {"type": "string"}, - "title": "LangGraphOutput", - } - - thread1 = {"configurable": {"thread_id": "1"}} - assert [ - c - async for c in graph.astream( - [0, 1], thread1, checkpoint_during=checkpoint_during + thread1 = {"configurable": {"thread_id": "1"}} + assert [ + c + async for c in graph.astream( + [0, 1], thread1, checkpoint_during=checkpoint_during + ) + ] == [ + {"submapper": "0"}, + {"mapper": "00"}, + {"submapper": "1"}, + {"mapper": "11"}, + { + "__interrupt__": ( + Interrupt( + value="question", + resumable=True, + ns=[AnyStr("graph:")], + when="during", + ), ) - ] == [ - {"submapper": "0"}, - {"mapper": "00"}, - {"submapper": "1"}, - {"mapper": "11"}, - { - "__interrupt__": ( - Interrupt( - value="question", - resumable=True, - ns=[AnyStr("graph:")], - when="during", - ), - ) - }, - ] + }, + ] - assert await graph.ainvoke( - Command(resume="answer"), thread1, checkpoint_during=checkpoint_during - ) == [ - "00answera", - "11answera", - ] + assert await graph.ainvoke( + Command(resume="answer"), thread1, checkpoint_during=checkpoint_during + ) == [ + "00answera", + "11answera", + ] @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_imp_task_cancel(checkpointer_name: str, checkpoint_during: bool) -> None: - async with awith_checkpointer(checkpointer_name) as checkpointer: - mapper_calls = 0 - mapper_cancels = 0 +async def test_imp_task_cancel( + async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool +) -> None: + mapper_calls = 0 + mapper_cancels = 0 - @task() - async def mapper(input: int) -> str: - nonlocal mapper_calls, mapper_cancels - mapper_calls += 1 - try: - await asyncio.sleep(1) - except asyncio.CancelledError: - mapper_cancels += 1 - raise - return str(input) * 2 + @task() + async def mapper(input: int) -> str: + nonlocal mapper_calls, mapper_cancels + mapper_calls += 1 + try: + await asyncio.sleep(1) + except asyncio.CancelledError: + mapper_cancels += 1 + raise + return str(input) * 2 - @entrypoint(checkpointer=checkpointer) - async def graph(input: list[int]) -> list[str]: - futures = [mapper(i) for i in input] - await asyncio.sleep(0.1) - futures.pop().cancel() # cancel one - mapped = await asyncio.gather(*futures) - answer = interrupt("question") - return [m + answer for m in mapped] + @entrypoint(checkpointer=async_checkpointer) + async def graph(input: list[int]) -> list[str]: + futures = [mapper(i) for i in input] + await asyncio.sleep(0.1) + futures.pop().cancel() # cancel one + mapped = await asyncio.gather(*futures) + answer = interrupt("question") + return [m + answer for m in mapped] - thread1 = {"configurable": {"thread_id": "1"}} - assert [ - c - async for c in graph.astream( - [0, 1], thread1, checkpoint_during=checkpoint_during + thread1 = {"configurable": {"thread_id": "1"}} + assert [ + c + async for c in graph.astream( + [0, 1], thread1, checkpoint_during=checkpoint_during + ) + ] == [ + {"mapper": "00"}, + { + "__interrupt__": ( + Interrupt( + value="question", + resumable=True, + ns=[AnyStr("graph:")], + when="during", + ), ) - ] == [ - {"mapper": "00"}, - { - "__interrupt__": ( - Interrupt( - value="question", - resumable=True, - ns=[AnyStr("graph:")], - when="during", - ), - ) - }, - ] - assert mapper_calls == 2 - assert mapper_cancels == 1 + }, + ] + assert mapper_calls == 2 + assert mapper_cancels == 1 - assert await graph.ainvoke( - Command(resume="answer"), thread1, checkpoint_during=checkpoint_during - ) == [ - "00answer", - ] - assert mapper_calls == 3 - assert mapper_cancels == 2 + assert await graph.ainvoke( + Command(resume="answer"), thread1, checkpoint_during=checkpoint_during + ) == [ + "00answer", + ] + assert mapper_calls == 3 + assert mapper_cancels == 2 @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_imp_sync_from_async( - checkpointer_name: str, checkpoint_during: bool + async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool ) -> None: - async with awith_checkpointer(checkpointer_name) as checkpointer: + @task() + def foo(state: dict) -> dict: + return {"a": state["a"] + "foo", "b": "bar"} - @task() - def foo(state: dict) -> dict: - return {"a": state["a"] + "foo", "b": "bar"} + @task + def bar(a: str, b: str, c: Optional[str] = None) -> dict: + return {"a": a + b, "c": (c or "") + "bark"} - @task - def bar(a: str, b: str, c: Optional[str] = None) -> dict: - return {"a": a + b, "c": (c or "") + "bark"} + @task() + def baz(state: dict) -> dict: + return {"a": state["a"] + "baz", "c": "something else"} - @task() - def baz(state: dict) -> dict: - return {"a": state["a"] + "baz", "c": "something else"} + @entrypoint(checkpointer=async_checkpointer) + def graph(state: dict) -> dict: + foo_result = foo(state).result() + fut_bar = bar(foo_result["a"], foo_result["b"]) + fut_baz = baz(fut_bar.result()) + return fut_baz.result() - @entrypoint(checkpointer=checkpointer) - def graph(state: dict) -> dict: - foo_result = foo(state).result() - fut_bar = bar(foo_result["a"], foo_result["b"]) - fut_baz = baz(fut_bar.result()) - return fut_baz.result() - - thread1 = {"configurable": {"thread_id": "1"}} - assert [ - c - async for c in graph.astream( - {"a": "0"}, thread1, checkpoint_during=checkpoint_during - ) - ] == [ - {"foo": {"a": "0foo", "b": "bar"}}, - {"bar": {"a": "0foobar", "c": "bark"}}, - {"baz": {"a": "0foobarbaz", "c": "something else"}}, - {"graph": {"a": "0foobarbaz", "c": "something else"}}, - ] + thread1 = {"configurable": {"thread_id": "1"}} + assert [ + c + async for c in graph.astream( + {"a": "0"}, thread1, checkpoint_during=checkpoint_during + ) + ] == [ + {"foo": {"a": "0foo", "b": "bar"}}, + {"bar": {"a": "0foobar", "c": "bark"}}, + {"baz": {"a": "0foobarbaz", "c": "something else"}}, + {"graph": {"a": "0foobarbaz", "c": "something else"}}, + ] @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_imp_stream_order( - checkpointer_name: str, checkpoint_during: bool + async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool ) -> None: - async with awith_checkpointer(checkpointer_name) as checkpointer: + @task() + async def foo(state: dict) -> dict: + return {"a": state["a"] + "foo", "b": "bar"} - @task() - async def foo(state: dict) -> dict: - return {"a": state["a"] + "foo", "b": "bar"} + @task + async def bar(a: str, b: str, c: Optional[str] = None) -> dict: + return {"a": a + b, "c": (c or "") + "bark"} - @task - async def bar(a: str, b: str, c: Optional[str] = None) -> dict: - return {"a": a + b, "c": (c or "") + "bark"} + @task() + async def baz(state: dict) -> dict: + return {"a": state["a"] + "baz", "c": "something else"} - @task() - async def baz(state: dict) -> dict: - return {"a": state["a"] + "baz", "c": "something else"} + @entrypoint(checkpointer=async_checkpointer) + async def graph(state: dict) -> dict: + foo_res = await foo(state) - @entrypoint(checkpointer=checkpointer) - async def graph(state: dict) -> dict: - foo_res = await foo(state) + fut_bar = bar(foo_res["a"], foo_res["b"]) + fut_baz = baz(await fut_bar) + return await fut_baz - fut_bar = bar(foo_res["a"], foo_res["b"]) - fut_baz = baz(await fut_bar) - return await fut_baz - - thread1 = {"configurable": {"thread_id": "1"}} - assert [ - c - async for c in graph.astream( - {"a": "0"}, thread1, checkpoint_during=checkpoint_during - ) - ] == [ - {"foo": {"a": "0foo", "b": "bar"}}, - {"bar": {"a": "0foobar", "c": "bark"}}, - {"baz": {"a": "0foobarbaz", "c": "something else"}}, - {"graph": {"a": "0foobarbaz", "c": "something else"}}, - ] + thread1 = {"configurable": {"thread_id": "1"}} + assert [ + c + async for c in graph.astream( + {"a": "0"}, thread1, checkpoint_during=checkpoint_during + ) + ] == [ + {"foo": {"a": "0foo", "b": "bar"}}, + {"bar": {"a": "0foobar", "c": "bark"}}, + {"baz": {"a": "0foobarbaz", "c": "something else"}}, + {"graph": {"a": "0foobarbaz", "c": "something else"}}, + ] async def test_send_dedupe_on_resume( @@ -3111,8 +3059,7 @@ async def test_send_dedupe_on_resume( assert history[1] == expected_history[2] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_send_react_interrupt(checkpointer_name: str) -> None: +async def test_send_react_interrupt(async_checkpointer: BaseCheckpointSaver) -> None: from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage ai_message = AIMessage( @@ -3166,55 +3113,75 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: } assert foo_called == 1 - async with awith_checkpointer(checkpointer_name) as checkpointer: - # simple interrupt-resume flow - foo_called = 0 - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"]) - thread1 = {"configurable": {"thread_id": "1"}} - assert await graph.ainvoke({"messages": [HumanMessage("hello")]}, thread1) == { - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ), - ] - } - assert foo_called == 0 - assert await graph.ainvoke(None, thread1) == { - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ), - _AnyIdToolMessage( - content="{'hi': [1, 2, 3]}", - tool_call_id=AnyStr(), - ), - ] - } - assert foo_called == 1 + # simple interrupt-resume flow + foo_called = 0 + graph = builder.compile(checkpointer=async_checkpointer, interrupt_before=["foo"]) + thread1 = {"configurable": {"thread_id": "1"}} + assert await graph.ainvoke({"messages": [HumanMessage("hello")]}, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + } + assert foo_called == 0 + assert await graph.ainvoke(None, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + _AnyIdToolMessage( + content="{'hi': [1, 2, 3]}", + tool_call_id=AnyStr(), + ), + ] + } + assert foo_called == 1 - # interrupt-update-resume flow - foo_called = 0 - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"]) - thread1 = {"configurable": {"thread_id": "2"}} - assert await graph.ainvoke({"messages": [HumanMessage("hello")]}, thread1) == { + # interrupt-update-resume flow + foo_called = 0 + graph = builder.compile(checkpointer=async_checkpointer, interrupt_before=["foo"]) + thread1 = {"configurable": {"thread_id": "2"}} + assert await graph.ainvoke({"messages": [HumanMessage("hello")]}, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + } + assert foo_called == 0 + + # get state should show the pending task + state = await graph.aget_state(thread1) + assert state == StateSnapshot( + values={ "messages": [ _AnyIdHumanMessage(content="hello"), _AnyIdAIMessage( @@ -3229,17 +3196,23 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: ], ), ] - } - assert foo_called == 0 - - # get state should show the pending task - state = await graph.aget_state(thread1) - assert state == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( + }, + next=("foo",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 1, + "source": "loop", + "writes": { + "agent": { + "messages": AIMessage( content="", + id="ai1", tool_calls=[ { "name": "foo", @@ -3248,129 +3221,123 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: "type": "tool_call", } ], - ), - ] + ) + } }, - next=("foo",), - config={ + "parents": {}, + "thread_id": "2", + }, + created_at=AnyStr(), + parent_config=( + { "configurable": { "thread_id": "2", "checkpoint_ns": "", "checkpoint_id": AnyStr(), } - }, - metadata={ - "step": 1, - "source": "loop", - "writes": { - "agent": { - "messages": AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ) - } - }, - "parents": {}, - "thread_id": "2", - }, - created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } + } + ), + tasks=( + PregelTask( + id=AnyStr(), + name="foo", + path=("__pregel_push", 0, False), + error=None, + interrupts=(), + state=None, + result=None, ), - tasks=( - PregelTask( - id=AnyStr(), - name="foo", - path=("__pregel_push", 0, False), - error=None, - interrupts=(), - state=None, - result=None, + ), + interrupts=(), + ) + + # remove the tool call, clearing the pending task + await graph.aupdate_state( + thread1, {"messages": AIMessage("Bye now", id=ai_message.id, tool_calls=[])} + ) + + # tool call no longer in pending tasks + assert await graph.aget_state(thread1) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="Bye now", + tool_calls=[], ), - ), - interrupts=(), - ) - - # remove the tool call, clearing the pending task - await graph.aupdate_state( - thread1, {"messages": AIMessage("Bye now", id=ai_message.id, tool_calls=[])} - ) - - # tool call no longer in pending tasks - assert await graph.aget_state(thread1) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( + ] + }, + next=(), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 2, + "source": "update", + "writes": { + "agent": { + "messages": _AnyIdAIMessage( content="Bye now", tool_calls=[], - ), - ] + ) + } }, - next=(), - config={ + "parents": {}, + "thread_id": "2", + }, + created_at=AnyStr(), + parent_config=( + { "configurable": { "thread_id": "2", "checkpoint_ns": "", "checkpoint_id": AnyStr(), } - }, - metadata={ - "step": 2, - "source": "update", - "writes": { - "agent": { - "messages": _AnyIdAIMessage( - content="Bye now", - tool_calls=[], - ) + } + ), + tasks=(), + interrupts=(), + ) + + # tool call not executed + assert await graph.ainvoke(None, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage(content="Bye now"), + ] + } + assert foo_called == 0 + + # interrupt-update-resume flow, creating new Send in update call + foo_called = 0 + graph = builder.compile(checkpointer=async_checkpointer, interrupt_before=["foo"]) + thread1 = {"configurable": {"thread_id": "3"}} + assert await graph.ainvoke({"messages": [HumanMessage("hello")]}, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", } - }, - "parents": {}, - "thread_id": "2", - }, - created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } + ], ), - tasks=(), - interrupts=(), - ) + ] + } + assert foo_called == 0 - # tool call not executed - assert await graph.ainvoke(None, thread1) == { - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage(content="Bye now"), - ] - } - assert foo_called == 0 - - # interrupt-update-resume flow, creating new Send in update call - foo_called = 0 - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"]) - thread1 = {"configurable": {"thread_id": "3"}} - assert await graph.ainvoke({"messages": [HumanMessage("hello")]}, thread1) == { + # get state should show the pending task + state = await graph.aget_state(thread1) + assert state == StateSnapshot( + values={ "messages": [ _AnyIdHumanMessage(content="hello"), _AnyIdAIMessage( @@ -3385,17 +3352,23 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: ], ), ] - } - assert foo_called == 0 - - # get state should show the pending task - state = await graph.aget_state(thread1) - assert state == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( + }, + next=("foo",), + config={ + "configurable": { + "thread_id": "3", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 1, + "source": "loop", + "writes": { + "agent": { + "messages": AIMessage( content="", + id="ai1", tool_calls=[ { "name": "foo", @@ -3404,70 +3377,62 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: "type": "tool_call", } ], - ), - ] + ) + } }, - next=("foo",), - config={ + "parents": {}, + "thread_id": "3", + }, + created_at=AnyStr(), + parent_config=( + { "configurable": { "thread_id": "3", "checkpoint_ns": "", "checkpoint_id": AnyStr(), } - }, - metadata={ - "step": 1, - "source": "loop", - "writes": { - "agent": { - "messages": AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ) - } - }, - "parents": {}, - "thread_id": "3", - }, - created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "3", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } + } + ), + tasks=( + PregelTask( + id=AnyStr(), + name="foo", + path=("__pregel_push", 0, False), + error=None, + interrupts=(), + state=None, + result=None, ), - tasks=( - PregelTask( - id=AnyStr(), - name="foo", - path=("__pregel_push", 0, False), - error=None, - interrupts=(), - state=None, - result=None, - ), - ), - interrupts=(), - ) + ), + interrupts=(), + ) - # replace the tool call, should clear previous send, create new one - await graph.aupdate_state( - thread1, - { - "messages": AIMessage( - "", - id=ai_message.id, + # replace the tool call, should clear previous send, create new one + await graph.aupdate_state( + thread1, + { + "messages": AIMessage( + "", + id=ai_message.id, + tool_calls=[ + { + "name": "foo", + "args": {"hi": [4, 5, 6]}, + "id": "tool1", + "type": "tool_call", + } + ], + ) + }, + ) + + # prev tool call no longer in pending tasks, new tool call is + assert await graph.aget_state(thread1) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", tool_calls=[ { "name": "foo", @@ -3476,16 +3441,23 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: "type": "tool_call", } ], - ) - }, - ) - - # prev tool call no longer in pending tasks, new tool call is - assert await graph.aget_state(thread1) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( + ), + ] + }, + next=("foo",), + config={ + "configurable": { + "thread_id": "3", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 2, + "source": "update", + "writes": { + "agent": { + "messages": _AnyIdAIMessage( content="", tool_calls=[ { @@ -3495,87 +3467,60 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: "type": "tool_call", } ], - ), - ] + ) + } }, - next=("foo",), - config={ + "parents": {}, + "thread_id": "3", + }, + created_at=AnyStr(), + parent_config=( + { "configurable": { "thread_id": "3", "checkpoint_ns": "", "checkpoint_id": AnyStr(), } - }, - metadata={ - "step": 2, - "source": "update", - "writes": { - "agent": { - "messages": _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [4, 5, 6]}, - "id": "tool1", - "type": "tool_call", - } - ], - ) - } - }, - "parents": {}, - "thread_id": "3", - }, - created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "3", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } + } + ), + tasks=( + PregelTask( + id=AnyStr(), + name="foo", + path=("__pregel_push", 0, False), + error=None, + interrupts=(), + state=None, + result=None, ), - tasks=( - PregelTask( - id=AnyStr(), - name="foo", - path=("__pregel_push", 0, False), - error=None, - interrupts=(), - state=None, - result=None, - ), + ), + interrupts=(), + ) + + # prev tool call not executed, new tool call is + assert await graph.ainvoke(None, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + AIMessage( + "", + id="ai1", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [4, 5, 6]}, + "id": "tool1", + "type": "tool_call", + } + ], ), - interrupts=(), - ) - - # prev tool call not executed, new tool call is - assert await graph.ainvoke(None, thread1) == { - "messages": [ - _AnyIdHumanMessage(content="hello"), - AIMessage( - "", - id="ai1", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [4, 5, 6]}, - "id": "tool1", - "type": "tool_call", - } - ], - ), - _AnyIdToolMessage(content="{'hi': [4, 5, 6]}", tool_call_id="tool1"), - ] - } - assert foo_called == 1 + _AnyIdToolMessage(content="{'hi': [4, 5, 6]}", tool_call_id="tool1"), + ] + } + assert foo_called == 1 -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_send_react_interrupt_control( - checkpointer_name: str, snapshot: SnapshotAssertion + async_checkpointer: BaseCheckpointSaver, snapshot: SnapshotAssertion ) -> None: from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage @@ -3603,7 +3548,7 @@ async def test_send_react_interrupt_control( builder.add_node(foo) builder.add_edge(START, "agent") graph = builder.compile() - if checkpointer_name == "memory": + if isinstance(async_checkpointer, InMemorySaver): assert graph.get_graph().draw_mermaid() == snapshot assert await graph.ainvoke({"messages": [HumanMessage("hello")]}) == { @@ -3628,55 +3573,75 @@ async def test_send_react_interrupt_control( } assert foo_called == 1 - async with awith_checkpointer(checkpointer_name) as checkpointer: - # simple interrupt-resume flow - foo_called = 0 - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"]) - thread1 = {"configurable": {"thread_id": "1"}} - assert await graph.ainvoke({"messages": [HumanMessage("hello")]}, thread1) == { - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ), - ] - } - assert foo_called == 0 - assert await graph.ainvoke(None, thread1) == { - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( - content="", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ), - _AnyIdToolMessage( - content="{'hi': [1, 2, 3]}", - tool_call_id=AnyStr(), - ), - ] - } - assert foo_called == 1 + # simple interrupt-resume flow + foo_called = 0 + graph = builder.compile(checkpointer=async_checkpointer, interrupt_before=["foo"]) + thread1 = {"configurable": {"thread_id": "1"}} + assert await graph.ainvoke({"messages": [HumanMessage("hello")]}, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + } + assert foo_called == 0 + assert await graph.ainvoke(None, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + _AnyIdToolMessage( + content="{'hi': [1, 2, 3]}", + tool_call_id=AnyStr(), + ), + ] + } + assert foo_called == 1 - # interrupt-update-resume flow - foo_called = 0 - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"]) - thread1 = {"configurable": {"thread_id": "2"}} - assert await graph.ainvoke({"messages": [HumanMessage("hello")]}, thread1) == { + # interrupt-update-resume flow + foo_called = 0 + graph = builder.compile(checkpointer=async_checkpointer, interrupt_before=["foo"]) + thread1 = {"configurable": {"thread_id": "2"}} + assert await graph.ainvoke({"messages": [HumanMessage("hello")]}, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + } + assert foo_called == 0 + + # get state should show the pending task + state = await graph.aget_state(thread1) + assert state == StateSnapshot( + values={ "messages": [ _AnyIdHumanMessage(content="hello"), _AnyIdAIMessage( @@ -3691,17 +3656,23 @@ async def test_send_react_interrupt_control( ], ), ] - } - assert foo_called == 0 - - # get state should show the pending task - state = await graph.aget_state(thread1) - assert state == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( + }, + next=("foo",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 1, + "source": "loop", + "writes": { + "agent": { + "messages": AIMessage( content="", + id="ai1", tool_calls=[ { "name": "foo", @@ -3710,127 +3681,99 @@ async def test_send_react_interrupt_control( "type": "tool_call", } ], - ), - ] + ) + } }, - next=("foo",), - config={ + "parents": {}, + "thread_id": "2", + }, + created_at=AnyStr(), + parent_config=( + { "configurable": { "thread_id": "2", "checkpoint_ns": "", "checkpoint_id": AnyStr(), } - }, - metadata={ - "step": 1, - "source": "loop", - "writes": { - "agent": { - "messages": AIMessage( - content="", - id="ai1", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ) - } - }, - "parents": {}, - "thread_id": "2", - }, - created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } + } + ), + tasks=( + PregelTask( + id=AnyStr(), + name="foo", + path=("__pregel_push", 0, False), + error=None, + interrupts=(), + state=None, + result=None, ), - tasks=( - PregelTask( - id=AnyStr(), - name="foo", - path=("__pregel_push", 0, False), - error=None, - interrupts=(), - state=None, - result=None, - ), - ), - interrupts=(), - ) + ), + interrupts=(), + ) - # remove the tool call, clearing the pending task - await graph.aupdate_state( - thread1, {"messages": AIMessage("Bye now", id=ai_message.id, tool_calls=[])} - ) + # remove the tool call, clearing the pending task + await graph.aupdate_state( + thread1, {"messages": AIMessage("Bye now", id=ai_message.id, tool_calls=[])} + ) - # tool call no longer in pending tasks - assert await graph.aget_state(thread1) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage( - content="Bye now", - tool_calls=[], - ), - ] - }, - next=(), - config={ - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "step": 2, - "source": "update", - "writes": { - "agent": { - "messages": _AnyIdAIMessage( - content="Bye now", - tool_calls=[], - ) - } - }, - "parents": {}, - "thread_id": "2", - }, - created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - tasks=(), - interrupts=(), - ) - - # tool call not executed - assert await graph.ainvoke(None, thread1) == { + # tool call no longer in pending tasks + assert await graph.aget_state(thread1) == StateSnapshot( + values={ "messages": [ _AnyIdHumanMessage(content="hello"), - _AnyIdAIMessage(content="Bye now"), + _AnyIdAIMessage( + content="Bye now", + tool_calls=[], + ), ] - } - assert foo_called == 0 + }, + next=(), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 2, + "source": "update", + "writes": { + "agent": { + "messages": _AnyIdAIMessage( + content="Bye now", + tool_calls=[], + ) + } + }, + "parents": {}, + "thread_id": "2", + }, + created_at=AnyStr(), + parent_config=( + { + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + } + ), + tasks=(), + interrupts=(), + ) + + # tool call not executed + assert await graph.ainvoke(None, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage(content="Bye now"), + ] + } + assert foo_called == 0 -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_max_concurrency(checkpointer_name: str) -> None: +async def test_max_concurrency(async_checkpointer: BaseCheckpointSaver) -> None: class Node: def __init__(self, name: str): self.name = name @@ -3882,18 +3825,16 @@ async def test_max_concurrency(checkpointer_name: str) -> None: assert node2.max_currently == 10 assert node2.currently == 0 - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["2"]) - thread1 = {"max_concurrency": 10, "configurable": {"thread_id": "1"}} + graph = builder.compile(checkpointer=async_checkpointer, interrupt_before=["2"]) + thread1 = {"max_concurrency": 10, "configurable": {"thread_id": "1"}} - assert await graph.ainvoke(["0"], thread1, debug=True) == ["0", "1"] - state = await graph.aget_state(thread1) - assert state.values == ["0", "1"] - assert await graph.ainvoke(None, thread1) == ["0", "1", *range(100), "3"] + assert await graph.ainvoke(["0"], thread1, debug=True) == ["0", "1"] + state = await graph.aget_state(thread1) + assert state.values == ["0", "1"] + assert await graph.ainvoke(None, thread1) == ["0", "1", *range(100), "3"] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_max_concurrency_control(checkpointer_name: str) -> None: +async def test_max_concurrency_control(async_checkpointer: BaseCheckpointSaver) -> None: async def node1(state) -> Command[Literal["2"]]: return Command(update=["1"], goto=[Send("2", idx) for idx in range(100)]) @@ -3920,7 +3861,7 @@ async def test_max_concurrency_control(checkpointer_name: str) -> None: builder.add_edge(START, "1") graph = builder.compile() - if checkpointer_name == "memory": + if isinstance(async_checkpointer, InMemorySaver): assert ( graph.get_graph().draw_mermaid() == """--- @@ -3958,17 +3899,15 @@ graph TD; assert node2_max_currently == 10 assert node2_currently == 0 - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["2"]) - thread1 = {"max_concurrency": 10, "configurable": {"thread_id": "1"}} + graph = builder.compile(checkpointer=async_checkpointer, interrupt_before=["2"]) + thread1 = {"max_concurrency": 10, "configurable": {"thread_id": "1"}} - assert await graph.ainvoke(["0"], thread1) == ["0", "1"] - assert await graph.ainvoke(None, thread1) == ["0", "1", *range(100), "3"] + assert await graph.ainvoke(["0"], thread1) == ["0", "1"] + assert await graph.ainvoke(None, thread1) == ["0", "1", *range(100), "3"] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_invoke_checkpoint_three( - mocker: MockerFixture, checkpointer_name: str + mocker: MockerFixture, async_checkpointer: BaseCheckpointSaver ) -> None: add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) @@ -3984,121 +3923,118 @@ async def test_invoke_checkpoint_three( | raise_if_above_10 ) - async with awith_checkpointer(checkpointer_name) as checkpointer: - app = Pregel( - nodes={"one": one}, - channels={ - "total": BinaryOperatorAggregate(int, operator.add), - "input": LastValue(int), - "output": LastValue(int), - }, - input_channels="input", - output_channels="output", - checkpointer=checkpointer, - debug=True, - ) + app = Pregel( + nodes={"one": one}, + channels={ + "total": BinaryOperatorAggregate(int, operator.add), + "input": LastValue(int), + "output": LastValue(int), + }, + input_channels="input", + output_channels="output", + checkpointer=async_checkpointer, + debug=True, + ) - thread_1 = {"configurable": {"thread_id": "1"}} - # total starts out as 0, so output is 0+2=2 - assert await app.ainvoke(2, thread_1) == 2 - state = await app.aget_state(thread_1) - assert state is not None - assert state.values.get("total") == 2 - assert ( - state.config["configurable"]["checkpoint_id"] - == (await checkpointer.aget(thread_1))["id"] - ) - # total is now 2, so output is 2+3=5 - assert await app.ainvoke(3, thread_1) == 5 - state = await app.aget_state(thread_1) - assert state is not None - assert state.values.get("total") == 7 - assert ( - state.config["configurable"]["checkpoint_id"] - == (await checkpointer.aget(thread_1))["id"] - ) - # total is now 2+5=7, so output would be 7+4=11, but raises ValueError - with pytest.raises(ValueError): - await app.ainvoke(4, thread_1) - # checkpoint is not updated - state = await app.aget_state(thread_1) - assert state is not None - assert state.values.get("total") == 7 - assert state.next == ("one",) - """we checkpoint inputs and it failed on "one", so the next node is one""" - # we can recover from error by sending new inputs - assert await app.ainvoke(2, thread_1) == 9 - state = await app.aget_state(thread_1) - assert state is not None - assert state.values.get("total") == 16, "total is now 7+9=16" - assert state.next == () + thread_1 = {"configurable": {"thread_id": "1"}} + # total starts out as 0, so output is 0+2=2 + assert await app.ainvoke(2, thread_1) == 2 + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 2 + assert ( + state.config["configurable"]["checkpoint_id"] + == (await async_checkpointer.aget(thread_1))["id"] + ) + # total is now 2, so output is 2+3=5 + assert await app.ainvoke(3, thread_1) == 5 + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 7 + assert ( + state.config["configurable"]["checkpoint_id"] + == (await async_checkpointer.aget(thread_1))["id"] + ) + # total is now 2+5=7, so output would be 7+4=11, but raises ValueError + with pytest.raises(ValueError): + await app.ainvoke(4, thread_1) + # checkpoint is not updated + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 7 + assert state.next == ("one",) + """we checkpoint inputs and it failed on "one", so the next node is one""" + # we can recover from error by sending new inputs + assert await app.ainvoke(2, thread_1) == 9 + state = await app.aget_state(thread_1) + assert state is not None + assert state.values.get("total") == 16, "total is now 7+9=16" + assert state.next == () - thread_2 = {"configurable": {"thread_id": "2"}} - # on a new thread, total starts out as 0, so output is 0+5=5 - assert await app.ainvoke(5, thread_2) == 5 - state = await app.aget_state({"configurable": {"thread_id": "1"}}) - assert state is not None - assert state.values.get("total") == 16 - assert state.next == () - state = await app.aget_state(thread_2) - assert state is not None - assert state.values.get("total") == 5 - assert state.next == () + thread_2 = {"configurable": {"thread_id": "2"}} + # on a new thread, total starts out as 0, so output is 0+5=5 + assert await app.ainvoke(5, thread_2) == 5 + state = await app.aget_state({"configurable": {"thread_id": "1"}}) + assert state is not None + assert state.values.get("total") == 16 + assert state.next == () + state = await app.aget_state(thread_2) + assert state is not None + assert state.values.get("total") == 5 + assert state.next == () - assert len([c async for c in app.aget_state_history(thread_1, limit=1)]) == 1 - # list all checkpoints for thread 1 - thread_1_history = [c async for c in app.aget_state_history(thread_1)] - # there are 7 checkpoints - assert len(thread_1_history) == 7 - assert Counter(c.metadata["source"] for c in thread_1_history) == { - "input": 4, - "loop": 3, - } - # sorted descending - assert ( - thread_1_history[0].config["configurable"]["checkpoint_id"] - > thread_1_history[1].config["configurable"]["checkpoint_id"] + assert len([c async for c in app.aget_state_history(thread_1, limit=1)]) == 1 + # list all checkpoints for thread 1 + thread_1_history = [c async for c in app.aget_state_history(thread_1)] + # there are 7 checkpoints + assert len(thread_1_history) == 7 + assert Counter(c.metadata["source"] for c in thread_1_history) == { + "input": 4, + "loop": 3, + } + # sorted descending + assert ( + thread_1_history[0].config["configurable"]["checkpoint_id"] + > thread_1_history[1].config["configurable"]["checkpoint_id"] + ) + # cursor pagination + cursored = [ + c + async for c in app.aget_state_history( + thread_1, limit=1, before=thread_1_history[0].config ) - # cursor pagination - cursored = [ - c - async for c in app.aget_state_history( - thread_1, limit=1, before=thread_1_history[0].config - ) - ] - assert len(cursored) == 1 - assert cursored[0].config == thread_1_history[1].config - # the last checkpoint - assert thread_1_history[0].values["total"] == 16 - # the first "loop" checkpoint - assert thread_1_history[-2].values["total"] == 2 - # can get each checkpoint using aget with config - assert (await checkpointer.aget(thread_1_history[0].config))[ - "id" - ] == thread_1_history[0].config["configurable"]["checkpoint_id"] - assert (await checkpointer.aget(thread_1_history[1].config))[ - "id" - ] == thread_1_history[1].config["configurable"]["checkpoint_id"] + ] + assert len(cursored) == 1 + assert cursored[0].config == thread_1_history[1].config + # the last checkpoint + assert thread_1_history[0].values["total"] == 16 + # the first "loop" checkpoint + assert thread_1_history[-2].values["total"] == 2 + # can get each checkpoint using aget with config + assert (await async_checkpointer.aget(thread_1_history[0].config))[ + "id" + ] == thread_1_history[0].config["configurable"]["checkpoint_id"] + assert (await async_checkpointer.aget(thread_1_history[1].config))[ + "id" + ] == thread_1_history[1].config["configurable"]["checkpoint_id"] - thread_1_next_config = await app.aupdate_state(thread_1_history[1].config, 10) - # update creates a new checkpoint - assert ( - thread_1_next_config["configurable"]["checkpoint_id"] - > thread_1_history[0].config["configurable"]["checkpoint_id"] - ) - # 1 more checkpoint in history - assert len([c async for c in app.aget_state_history(thread_1)]) == 8 - assert Counter( - [c.metadata["source"] async for c in app.aget_state_history(thread_1)] - ) == { - "update": 1, - "input": 4, - "loop": 3, - } - # the latest checkpoint is the updated one - assert await app.aget_state(thread_1) == await app.aget_state( - thread_1_next_config - ) + thread_1_next_config = await app.aupdate_state(thread_1_history[1].config, 10) + # update creates a new checkpoint + assert ( + thread_1_next_config["configurable"]["checkpoint_id"] + > thread_1_history[0].config["configurable"]["checkpoint_id"] + ) + # 1 more checkpoint in history + assert len([c async for c in app.aget_state_history(thread_1)]) == 8 + assert Counter( + [c.metadata["source"] async for c in app.aget_state_history(thread_1)] + ) == { + "update": 1, + "input": 4, + "loop": 3, + } + # the latest checkpoint is the updated one + assert await app.aget_state(thread_1) == await app.aget_state(thread_1_next_config) async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None: @@ -4137,9 +4073,8 @@ async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) - ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_invoke_join_then_call_other_pregel( - mocker: MockerFixture, checkpointer_name: str + mocker: MockerFixture, async_checkpointer: BaseCheckpointSaver ) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x]) @@ -4193,16 +4128,15 @@ async def test_invoke_join_then_call_other_pregel( 27 for _ in range(10) ] - async with awith_checkpointer(checkpointer_name) as checkpointer: - # add checkpointer - app.checkpointer = checkpointer - # subgraph is called twice, and that works - assert await app.ainvoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27 + # add checkpointer + app.checkpointer = async_checkpointer + # subgraph is called twice, and that works + assert await app.ainvoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27 - # set inner graph checkpointer NeverCheckpoint - inner_app.checkpointer = False - # subgraph still called twice, but checkpointing for inner graph is disabled - assert await app.ainvoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27 + # set inner graph checkpointer NeverCheckpoint + inner_app.checkpointer = False + # subgraph still called twice, but checkpointing for inner graph is disabled + assert await app.ainvoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27 async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None: @@ -4335,8 +4269,9 @@ async def test_conditional_entrypoint_graph_state() -> None: ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_in_one_fan_out_state_graph_waiting_edge(checkpointer_name: str) -> None: +async def test_in_one_fan_out_state_graph_waiting_edge( + async_checkpointer: BaseCheckpointSaver, +) -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] ) -> list[str]: @@ -4398,35 +4333,33 @@ async def test_in_one_fan_out_state_graph_waiting_edge(checkpointer_name: str) - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] - async with awith_checkpointer(checkpointer_name) as checkpointer: - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["retriever_one"], + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_after=["retriever_one"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf"}, config ) - config = {"configurable": {"thread_id": "1"}} + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ] - assert [ - c - async for c in app_w_interrupt.astream( - {"query": "what is weather in sf"}, config - ) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - {"__interrupt__": ()}, - ] - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] @pytest.mark.parametrize("use_waiting_edge", (True, False)) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_in_one_fan_out_state_graph_defer_node( - checkpointer_name: str, use_waiting_edge: bool + async_checkpointer: BaseCheckpointSaver, use_waiting_edge: bool ) -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] @@ -4493,35 +4426,33 @@ async def test_in_one_fan_out_state_graph_defer_node( {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] - async with awith_checkpointer(checkpointer_name) as checkpointer: - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["retriever_one"], + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_after=["retriever_one"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf"}, config ) - config = {"configurable": {"thread_id": "1"}} + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ] - assert [ - c - async for c in app_w_interrupt.astream( - {"query": "what is weather in sf"}, config - ) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - {"__interrupt__": ()}, - ] - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] @pytest.mark.parametrize("with_path_map", (True, False)) -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_in_one_fan_out_state_graph_then_defer_node( - checkpointer_name: str, with_path_map: bool + async_checkpointer: BaseCheckpointSaver, with_path_map: bool ) -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] @@ -4586,34 +4517,32 @@ async def test_in_one_fan_out_state_graph_then_defer_node( {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] - async with awith_checkpointer(checkpointer_name) as checkpointer: - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["retriever_one"], + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_after=["retriever_one"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf"}, config ) - config = {"configurable": {"thread_id": "1"}} + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ] - assert [ - c - async for c in app_w_interrupt.astream( - {"query": "what is weather in sf"}, config - ) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - {"__interrupt__": ()}, - ] - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_in_one_fan_out_state_graph_waiting_edge_via_branch( - snapshot: SnapshotAssertion, checkpointer_name: str + async_checkpointer: BaseCheckpointSaver, ) -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] @@ -4678,29 +4607,28 @@ async def test_in_one_fan_out_state_graph_waiting_edge_via_branch( {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] - async with awith_checkpointer(checkpointer_name) as checkpointer: - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["retriever_one"], + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_after=["retriever_one"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf"}, config ) - config = {"configurable": {"thread_id": "1"}} + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ] - assert [ - c - async for c in app_w_interrupt.astream( - {"query": "what is weather in sf"}, config - ) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - {"__interrupt__": ()}, - ] - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] async def test_nested_pydantic_models() -> None: @@ -4828,9 +4756,8 @@ async def test_nested_pydantic_models() -> None: assert result == {**inputs, **update} -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( - checkpointer_name: str, + async_checkpointer: BaseCheckpointSaver, ) -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] @@ -4915,79 +4842,77 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] - async with awith_checkpointer(checkpointer_name) as checkpointer: - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["retriever_one"], + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_after=["retriever_one"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf"}, config ) - config = {"configurable": {"thread_id": "1"}} + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ] - assert [ - c - async for c in app_w_interrupt.astream( - {"query": "what is weather in sf"}, config - ) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - {"__interrupt__": ()}, - ] + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] - - assert await app_w_interrupt.aget_state(config) == StateSnapshot( - values={ - "query": "analyzed: query: what is weather in sf", - "answer": "doc1,doc2,doc3,doc4", - "docs": ["doc1", "doc2", "doc3", "doc4"], - }, - tasks=(), - next=(), - config={ + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "query": "analyzed: query: what is weather in sf", + "answer": "doc1,doc2,doc3,doc4", + "docs": ["doc1", "doc2", "doc3", "doc4"], + }, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + "step": 4, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": "", "checkpoint_id": AnyStr(), } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - "step": 4, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - interrupts=(), - ) - - assert await app_w_interrupt.aupdate_state( - config, {"docs": ["doc5"]}, as_node="rewrite_query" - ) == { - "configurable": { - "thread_id": "1", - "checkpoint_id": AnyStr(), - "checkpoint_ns": "", } + ), + interrupts=(), + ) + + assert await app_w_interrupt.aupdate_state( + config, {"docs": ["doc5"]}, as_node="rewrite_query" + ) == { + "configurable": { + "thread_id": "1", + "checkpoint_id": AnyStr(), + "checkpoint_ns": "", } + } -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2( - snapshot: SnapshotAssertion, checkpointer_name: str + snapshot: SnapshotAssertion, async_checkpointer: BaseCheckpointSaver ) -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] @@ -5051,7 +4976,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydant app = workflow.compile() - if SHOULD_CHECK_SNAPSHOTS and checkpointer_name == "memory": + if isinstance(async_checkpointer, InMemorySaver): assert app.get_graph().draw_mermaid(with_styles=False) == snapshot assert app.get_input_schema().model_json_schema() == snapshot assert app.get_output_schema().model_json_schema() == snapshot @@ -5081,44 +5006,42 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydant {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] - async with awith_checkpointer(checkpointer_name) as checkpointer: - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["retriever_one"], + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_after=["retriever_one"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf", "inner": {"yo": 1}}, config ) - config = {"configurable": {"thread_id": "1"}} + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ] - assert [ - c - async for c in app_w_interrupt.astream( - {"query": "what is weather in sf", "inner": {"yo": 1}}, config - ) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - {"__interrupt__": ()}, - ] + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] - - assert await app_w_interrupt.aupdate_state( - config, {"docs": ["doc5"]}, as_node="rewrite_query" - ) == { - "configurable": { - "thread_id": "1", - "checkpoint_id": AnyStr(), - "checkpoint_ns": "", - } + assert await app_w_interrupt.aupdate_state( + config, {"docs": ["doc5"]}, as_node="rewrite_query" + ) == { + "configurable": { + "thread_id": "1", + "checkpoint_id": AnyStr(), + "checkpoint_ns": "", } + } -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( - checkpointer_name: str, + async_checkpointer: BaseCheckpointSaver, ) -> None: def sorted_add( x: list[str], y: Union[list[str], list[tuple[str, str]]] @@ -5187,30 +5110,29 @@ async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] - async with awith_checkpointer(checkpointer_name) as checkpointer: - app_w_interrupt = workflow.compile( - checkpointer=checkpointer, - interrupt_after=["retriever_one"], + app_w_interrupt = workflow.compile( + checkpointer=async_checkpointer, + interrupt_after=["retriever_one"], + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + {"query": "what is weather in sf"}, config ) - config = {"configurable": {"thread_id": "1"}} + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"qa": {"answer": ""}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ] - assert [ - c - async for c in app_w_interrupt.astream( - {"query": "what is weather in sf"}, config - ) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"qa": {"answer": ""}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - {"__interrupt__": ()}, - ] - - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - ] + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] @pytest.mark.parametrize("with_cache", [True, False]) @@ -5671,8 +5593,9 @@ async def test_subgraph_checkpoint_true_interrupt( ) == {"foo": "hi! foobaz"} -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_stream_subgraphs_during_execution(checkpointer_name: str) -> None: +async def test_stream_subgraphs_during_execution( + async_checkpointer: BaseCheckpointSaver, +) -> None: class InnerState(TypedDict): my_key: Annotated[str, operator.add] my_other_key: str @@ -5714,42 +5637,42 @@ async def test_stream_subgraphs_during_execution(checkpointer_name: str) -> None graph.add_edge(["inner", "outer_1"], "outer_2") graph.add_edge("outer_2", END) - async with awith_checkpointer(checkpointer_name) as checkpointer: - app = graph.compile(checkpointer=checkpointer) + app = graph.compile(checkpointer=async_checkpointer) - start = perf_counter() - chunks: list[tuple[float, Any]] = [] - config = {"configurable": {"thread_id": "2"}} - async for c in app.astream({"my_key": ""}, config, subgraphs=True): - chunks.append((round(perf_counter() - start, 1), c)) - for idx in range(len(chunks)): - elapsed, c = chunks[idx] - chunks[idx] = (round(elapsed - chunks[0][0], 1), c) + start = perf_counter() + chunks: list[tuple[float, Any]] = [] + config = {"configurable": {"thread_id": "2"}} + async for c in app.astream({"my_key": ""}, config, subgraphs=True): + chunks.append((round(perf_counter() - start, 1), c)) + for idx in range(len(chunks)): + elapsed, c = chunks[idx] + chunks[idx] = (round(elapsed - chunks[0][0], 1), c) - assert chunks == [ - # arrives before "inner" finishes + assert chunks == [ + # arrives before "inner" finishes + ( + FloatBetween(0.0, 0.1), ( - FloatBetween(0.0, 0.1), - ( - (AnyStr("inner:"),), - {"inner_1": {"my_key": "got here", "my_other_key": ""}}, - ), + (AnyStr("inner:"),), + {"inner_1": {"my_key": "got here", "my_other_key": ""}}, ), - (FloatBetween(0.2, 0.4), ((), {"outer_1": {"my_key": " and parallel"}})), + ), + (FloatBetween(0.2, 0.4), ((), {"outer_1": {"my_key": " and parallel"}})), + ( + FloatBetween(0.5, 0.8), ( - FloatBetween(0.5, 0.8), - ( - (AnyStr("inner:"),), - {"inner_2": {"my_key": " and there", "my_other_key": "got here"}}, - ), + (AnyStr("inner:"),), + {"inner_2": {"my_key": " and there", "my_other_key": "got here"}}, ), - (FloatBetween(0.5, 0.8), ((), {"inner": {"my_key": "got here and there"}})), - (FloatBetween(0.5, 0.8), ((), {"outer_2": {"my_key": " and back again"}})), - ] + ), + (FloatBetween(0.5, 0.8), ((), {"inner": {"my_key": "got here and there"}})), + (FloatBetween(0.5, 0.8), ((), {"outer_2": {"my_key": " and back again"}})), + ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_stream_buffering_single_node(checkpointer_name: str) -> None: +async def test_stream_buffering_single_node( + async_checkpointer: BaseCheckpointSaver, +) -> None: class State(TypedDict): my_key: Annotated[str, operator.add] @@ -5764,24 +5687,22 @@ async def test_stream_buffering_single_node(checkpointer_name: str) -> None: builder.add_edge(START, "node") builder.add_edge("node", END) - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) + graph = builder.compile(checkpointer=async_checkpointer) - start = perf_counter() - chunks: list[tuple[float, Any]] = [] - config = {"configurable": {"thread_id": "2"}} - async for c in graph.astream({"my_key": ""}, config, stream_mode="custom"): - chunks.append((round(perf_counter() - start, 1), c)) + start = perf_counter() + chunks: list[tuple[float, Any]] = [] + config = {"configurable": {"thread_id": "2"}} + async for c in graph.astream({"my_key": ""}, config, stream_mode="custom"): + chunks.append((round(perf_counter() - start, 1), c)) - assert chunks == [ - (FloatBetween(0.0, 0.1), "Before sleep"), - (FloatBetween(0.2, 0.3), "After sleep"), - ] + assert chunks == [ + (FloatBetween(0.0, 0.1), "Before sleep"), + (FloatBetween(0.2, 0.3), "After sleep"), + ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_nested_graph_interrupts_parallel( - checkpointer_name: str, checkpoint_during: bool + async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool ) -> None: class InnerState(TypedDict): my_key: Annotated[str, operator.add] @@ -5826,152 +5747,147 @@ async def test_nested_graph_interrupts_parallel( graph.add_edge(["inner", "outer_1"], "outer_2") graph.set_finish_point("outer_2") - async with awith_checkpointer(checkpointer_name) as checkpointer: - app = graph.compile(checkpointer=checkpointer) + app = graph.compile(checkpointer=async_checkpointer) - # test invoke w/ nested interrupt - config = {"configurable": {"thread_id": "1"}} - assert await app.ainvoke( - {"my_key": ""}, config, checkpoint_during=checkpoint_during - ) == { - "my_key": " and parallel", - } + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert await app.ainvoke( + {"my_key": ""}, config, checkpoint_during=checkpoint_during + ) == { + "my_key": " and parallel", + } - assert await app.ainvoke(None, config, checkpoint_during=checkpoint_during) == { - "my_key": "got here and there and parallel and back again", - } + assert await app.ainvoke(None, config, checkpoint_during=checkpoint_during) == { + "my_key": "got here and there and parallel and back again", + } - # below combo of assertions is asserting two things - # - outer_1 finishes before inner interrupts (because we see its output in stream, which only happens after node finishes) - # - the writes of outer are persisted in 1st call and used in 2nd call, ie outer isn't called again (because we dont see outer_1 output again in 2nd stream) - # test stream updates w/ nested interrupt - config = {"configurable": {"thread_id": "2"}} - assert [ - c - async for c in app.astream( - {"my_key": ""}, - config, - subgraphs=True, - checkpoint_during=checkpoint_during, - ) - ] == [ - # we got to parallel node first - ((), {"outer_1": {"my_key": " and parallel"}}), - ( - (AnyStr("inner:"),), - {"inner_1": {"my_key": "got here", "my_other_key": ""}}, - ), - ((), {"__interrupt__": ()}), - ] - assert [ - c - async for c in app.astream( - None, config, checkpoint_during=checkpoint_during - ) - ] == [ - {"outer_1": {"my_key": " and parallel"}, "__metadata__": {"cached": True}}, - {"inner": {"my_key": "got here and there"}}, - {"outer_2": {"my_key": " and back again"}}, - ] - - # test stream values w/ nested interrupt - config = {"configurable": {"thread_id": "3"}} - assert [ - c - async for c in app.astream( - {"my_key": ""}, - config, - stream_mode="values", - checkpoint_during=checkpoint_during, - ) - ] == [ + # below combo of assertions is asserting two things + # - outer_1 finishes before inner interrupts (because we see its output in stream, which only happens after node finishes) + # - the writes of outer are persisted in 1st call and used in 2nd call, ie outer isn't called again (because we dont see outer_1 output again in 2nd stream) + # test stream updates w/ nested interrupt + config = {"configurable": {"thread_id": "2"}} + assert [ + c + async for c in app.astream( {"my_key": ""}, - {"my_key": " and parallel"}, - ] - assert [ - c - async for c in app.astream( - None, config, stream_mode="values", checkpoint_during=checkpoint_during - ) - ] == [ - {"my_key": ""}, - {"my_key": "got here and there and parallel"}, - {"my_key": "got here and there and parallel and back again"}, - ] + config, + subgraphs=True, + checkpoint_during=checkpoint_during, + ) + ] == [ + # we got to parallel node first + ((), {"outer_1": {"my_key": " and parallel"}}), + ( + (AnyStr("inner:"),), + {"inner_1": {"my_key": "got here", "my_other_key": ""}}, + ), + ((), {"__interrupt__": ()}), + ] + assert [ + c async for c in app.astream(None, config, checkpoint_during=checkpoint_during) + ] == [ + {"outer_1": {"my_key": " and parallel"}, "__metadata__": {"cached": True}}, + {"inner": {"my_key": "got here and there"}}, + {"outer_2": {"my_key": " and back again"}}, + ] - # # test interrupts BEFORE the parallel node - app = graph.compile(checkpointer=checkpointer, interrupt_before=["outer_1"]) - config = {"configurable": {"thread_id": "4"}} - assert [ - c - async for c in app.astream( - {"my_key": ""}, - config, - stream_mode="values", - checkpoint_during=checkpoint_during, - ) - ] == [ + # test stream values w/ nested interrupt + config = {"configurable": {"thread_id": "3"}} + assert [ + c + async for c in app.astream( {"my_key": ""}, - ] - # while we're waiting for the node w/ interrupt inside to finish - assert [ - c - async for c in app.astream( - None, config, stream_mode="values", checkpoint_during=checkpoint_during - ) - ] == [ - {"my_key": ""}, - {"my_key": " and parallel"}, - ] - assert [ - c - async for c in app.astream( - None, config, stream_mode="values", checkpoint_during=checkpoint_during - ) - ] == [ - {"my_key": ""}, - {"my_key": "got here and there and parallel"}, - {"my_key": "got here and there and parallel and back again"}, - ] + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) + ] == [ + {"my_key": ""}, + {"my_key": " and parallel"}, + ] + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ + {"my_key": ""}, + {"my_key": "got here and there and parallel"}, + {"my_key": "got here and there and parallel and back again"}, + ] - # test interrupts AFTER the parallel node - app = graph.compile(checkpointer=checkpointer, interrupt_after=["outer_1"]) - config = {"configurable": {"thread_id": "5"}} - assert [ - c - async for c in app.astream( - {"my_key": ""}, - config, - stream_mode="values", - checkpoint_during=checkpoint_during, - ) - ] == [ + # # test interrupts BEFORE the parallel node + app = graph.compile(checkpointer=async_checkpointer, interrupt_before=["outer_1"]) + config = {"configurable": {"thread_id": "4"}} + assert [ + c + async for c in app.astream( {"my_key": ""}, - {"my_key": " and parallel"}, - ] - assert [ - c - async for c in app.astream( - None, config, stream_mode="values", checkpoint_during=checkpoint_during - ) - ] == [ + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) + ] == [ + {"my_key": ""}, + ] + # while we're waiting for the node w/ interrupt inside to finish + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ + {"my_key": ""}, + {"my_key": " and parallel"}, + ] + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ + {"my_key": ""}, + {"my_key": "got here and there and parallel"}, + {"my_key": "got here and there and parallel and back again"}, + ] + + # test interrupts AFTER the parallel node + app = graph.compile(checkpointer=async_checkpointer, interrupt_after=["outer_1"]) + config = {"configurable": {"thread_id": "5"}} + assert [ + c + async for c in app.astream( {"my_key": ""}, - {"my_key": "got here and there and parallel"}, - ] - assert [ - c - async for c in app.astream( - None, config, stream_mode="values", checkpoint_during=checkpoint_during - ) - ] == [ - {"my_key": "got here and there and parallel"}, - {"my_key": "got here and there and parallel and back again"}, - ] + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) + ] == [ + {"my_key": ""}, + {"my_key": " and parallel"}, + ] + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ + {"my_key": ""}, + {"my_key": "got here and there and parallel"}, + ] + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ + {"my_key": "got here and there and parallel"}, + {"my_key": "got here and there and parallel and back again"}, + ] -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_doubly_nested_graph_interrupts( - checkpointer_name: str, checkpoint_during: bool + async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool ) -> None: class State(TypedDict): my_key: str @@ -6020,78 +5936,74 @@ async def test_doubly_nested_graph_interrupts( graph.add_edge("child", "parent_2") graph.set_finish_point("parent_2") - async with awith_checkpointer(checkpointer_name) as checkpointer: - app = graph.compile(checkpointer=checkpointer) + app = graph.compile(checkpointer=async_checkpointer) - # test invoke w/ nested interrupt - config = {"configurable": {"thread_id": "1"}} - assert await app.ainvoke( + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + assert await app.ainvoke( + {"my_key": "my value"}, config, checkpoint_during=checkpoint_during + ) == { + "my_key": "hi my value", + } + + assert await app.ainvoke(None, config, checkpoint_during=checkpoint_during) == { + "my_key": "hi my value here and there and back again", + } + + # test stream updates w/ nested interrupt + nodes: list[str] = [] + config = { + "configurable": {"thread_id": "2", CONFIG_KEY_NODE_FINISHED: nodes.append} + } + assert [ + c + async for c in app.astream( {"my_key": "my value"}, config, checkpoint_during=checkpoint_during - ) == { - "my_key": "hi my value", - } + ) + ] == [ + {"parent_1": {"my_key": "hi my value"}}, + {"__interrupt__": ()}, + ] + assert nodes == ["parent_1", "grandchild_1"] + assert [ + c async for c in app.astream(None, config, checkpoint_during=checkpoint_during) + ] == [ + {"child": {"my_key": "hi my value here and there"}}, + {"parent_2": {"my_key": "hi my value here and there and back again"}}, + ] + assert nodes == [ + "parent_1", + "grandchild_1", + "grandchild_2", + "child_1", + "child", + "parent_2", + ] - assert await app.ainvoke(None, config, checkpoint_during=checkpoint_during) == { - "my_key": "hi my value here and there and back again", - } - - # test stream updates w/ nested interrupt - nodes: list[str] = [] - config = { - "configurable": {"thread_id": "2", CONFIG_KEY_NODE_FINISHED: nodes.append} - } - assert [ - c - async for c in app.astream( - {"my_key": "my value"}, config, checkpoint_during=checkpoint_during - ) - ] == [ - {"parent_1": {"my_key": "hi my value"}}, - {"__interrupt__": ()}, - ] - assert nodes == ["parent_1", "grandchild_1"] - assert [ - c - async for c in app.astream( - None, config, checkpoint_during=checkpoint_during - ) - ] == [ - {"child": {"my_key": "hi my value here and there"}}, - {"parent_2": {"my_key": "hi my value here and there and back again"}}, - ] - assert nodes == [ - "parent_1", - "grandchild_1", - "grandchild_2", - "child_1", - "child", - "parent_2", - ] - - # test stream values w/ nested interrupt - config = {"configurable": {"thread_id": "3"}} - assert [ - c - async for c in app.astream( - {"my_key": "my value"}, - config, - stream_mode="values", - checkpoint_during=checkpoint_during, - ) - ] == [ + # test stream values w/ nested interrupt + config = {"configurable": {"thread_id": "3"}} + assert [ + c + async for c in app.astream( {"my_key": "my value"}, - {"my_key": "hi my value"}, - ] - assert [ - c - async for c in app.astream( - None, config, stream_mode="values", checkpoint_during=checkpoint_during - ) - ] == [ - {"my_key": "hi my value"}, - {"my_key": "hi my value here and there"}, - {"my_key": "hi my value here and there and back again"}, - ] + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) + ] == [ + {"my_key": "my value"}, + {"my_key": "hi my value"}, + ] + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ + {"my_key": "hi my value"}, + {"my_key": "hi my value here and there"}, + {"my_key": "hi my value here and there and back again"}, + ] async def test_checkpoint_metadata(async_checkpointer: BaseCheckpointSaver) -> None: @@ -6273,9 +6185,8 @@ async def test_checkpointer_null_pending_writes() -> None: ] * 4 -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_store_injected_async( - checkpointer_name: str, async_store: BaseStore + async_checkpointer: BaseCheckpointSaver, async_store: BaseStore ) -> None: class State(TypedDict): count: Annotated[int, operator.add] @@ -6332,46 +6243,43 @@ async def test_store_injected_async( builder.add_node(f"node_{i}", Node(i)) builder.add_edge("__start__", f"node_{i}") - async with ( - awith_checkpointer(checkpointer_name) as checkpointer, - ): - graph = builder.compile(store=async_store, checkpointer=checkpointer) + graph = builder.compile(store=async_store, checkpointer=async_checkpointer) - # Test batch operations with multiple threads - results = await graph.abatch( - [{"count": 0}] * M, - ([{"configurable": {"thread_id": str(uuid.uuid4())}}] * (M - 1)) - + [{"configurable": {"thread_id": thread_1}}], - ) - result = results[-1] - assert result == {"count": N + 1} - returned_doc = (await async_store.aget(namespace, doc_id)).value - assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 0} - assert len(await async_store.asearch(namespace)) == 1 + # Test batch operations with multiple threads + results = await graph.abatch( + [{"count": 0}] * M, + ([{"configurable": {"thread_id": str(uuid.uuid4())}}] * (M - 1)) + + [{"configurable": {"thread_id": thread_1}}], + ) + result = results[-1] + assert result == {"count": N + 1} + returned_doc = (await async_store.aget(namespace, doc_id)).value + assert returned_doc == {**doc, "from_thread": thread_1, "some_val": 0} + assert len(await async_store.asearch(namespace)) == 1 - # Check results after another turn of the same thread - result = await graph.ainvoke( - {"count": 0}, {"configurable": {"thread_id": thread_1}} - ) - assert result == {"count": (N + 1) * 2} - returned_doc = (await async_store.aget(namespace, doc_id)).value - assert returned_doc == {**doc, "from_thread": thread_1, "some_val": N + 1} - assert len(await async_store.asearch(namespace)) == 1 + # Check results after another turn of the same thread + result = await graph.ainvoke( + {"count": 0}, {"configurable": {"thread_id": thread_1}} + ) + assert result == {"count": (N + 1) * 2} + returned_doc = (await async_store.aget(namespace, doc_id)).value + assert returned_doc == {**doc, "from_thread": thread_1, "some_val": N + 1} + assert len(await async_store.asearch(namespace)) == 1 - # Test with a different thread - result = await graph.ainvoke( - {"count": 0}, {"configurable": {"thread_id": thread_2}} - ) - assert result == {"count": N + 1} - returned_doc = (await async_store.aget(namespace, doc_id)).value - assert returned_doc == { - **doc, - "from_thread": thread_2, - "some_val": 0, - } # Overwrites the whole doc - assert ( - len(await async_store.asearch(namespace)) == 1 - ) # still overwriting the same one + # Test with a different thread + result = await graph.ainvoke( + {"count": 0}, {"configurable": {"thread_id": thread_2}} + ) + assert result == {"count": N + 1} + returned_doc = (await async_store.aget(namespace, doc_id)).value + assert returned_doc == { + **doc, + "from_thread": thread_2, + "some_val": 0, + } # Overwrites the whole doc + assert ( + len(await async_store.asearch(namespace)) == 1 + ) # still overwriting the same one async def test_debug_retry(async_checkpointer: BaseCheckpointSaver): @@ -6610,8 +6518,7 @@ async def test_debug_nested_subgraphs(async_checkpointer: BaseCheckpointSaver): assert stream_task.get("state") == history_task.state -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_parent_command(checkpointer_name: str) -> None: +async def test_parent_command(async_checkpointer: BaseCheckpointSaver) -> None: from langchain_core.messages import BaseMessage from langchain_core.tools import tool @@ -6633,69 +6540,66 @@ async def test_parent_command(checkpointer_name: str) -> None: builder = StateGraph(CustomParentState) builder.add_node("alice", subgraph) builder.add_edge(START, "alice") - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) - config = {"configurable": {"thread_id": "1"}} + graph = builder.compile(checkpointer=async_checkpointer) - assert await graph.ainvoke( - {"messages": [("user", "get user name")]}, config - ) == { + config = {"configurable": {"thread_id": "1"}} + + assert await graph.ainvoke({"messages": [("user", "get user name")]}, config) == { + "messages": [ + _AnyIdHumanMessage( + content="get user name", additional_kwargs={}, response_metadata={} + ), + ], + "user_name": "Meow", + } + assert await graph.aget_state(config) == StateSnapshot( + values={ "messages": [ _AnyIdHumanMessage( - content="get user name", additional_kwargs={}, response_metadata={} + content="get user name", + additional_kwargs={}, + response_metadata={}, ), ], "user_name": "Meow", - } - assert await graph.aget_state(config) == StateSnapshot( - values={ - "messages": [ - _AnyIdHumanMessage( - content="get user name", - additional_kwargs={}, - response_metadata={}, - ), - ], - "user_name": "Meow", + }, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "alice": { + "user_name": "Meow", + } }, - next=(), - config={ + "thread_id": "1", + "step": 1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config=( + { "configurable": { "thread_id": "1", "checkpoint_ns": "", "checkpoint_id": AnyStr(), } - }, - metadata={ - "source": "loop", - "writes": { - "alice": { - "user_name": "Meow", - } - }, - "thread_id": "1", - "step": 1, - "parents": {}, - }, - created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - tasks=(), - interrupts=(), - ) + } + ), + tasks=(), + interrupts=(), + ) @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_interrupt_subgraph(checkpointer_name: str): +async def test_interrupt_subgraph(async_checkpointer: BaseCheckpointSaver) -> None: class State(TypedDict): baz: str @@ -6716,19 +6620,17 @@ async def test_interrupt_subgraph(checkpointer_name: str): builder.add_edge(START, "foo") builder.add_edge("foo", "bar") - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) + graph = builder.compile(checkpointer=async_checkpointer) - thread1 = {"configurable": {"thread_id": "1"}} - # First run, interrupted at bar - assert await graph.ainvoke({"baz": ""}, thread1) - # Resume with answer - assert await graph.ainvoke(Command(resume="bar"), thread1) + thread1 = {"configurable": {"thread_id": "1"}} + # First run, interrupted at bar + assert await graph.ainvoke({"baz": ""}, thread1) + # Resume with answer + assert await graph.ainvoke(Command(resume="bar"), thread1) @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_interrupt_multiple(checkpointer_name: str): +async def test_interrupt_multiple(async_checkpointer: BaseCheckpointSaver): class State(TypedDict): my_key: Annotated[str, operator.add] @@ -6741,58 +6643,56 @@ async def test_interrupt_multiple(checkpointer_name: str): builder.add_node("node", node) builder.add_edge(START, "node") - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) - thread1 = {"configurable": {"thread_id": "1"}} + graph = builder.compile(checkpointer=async_checkpointer) + thread1 = {"configurable": {"thread_id": "1"}} - assert [ - e async for e in graph.astream({"my_key": "DE", "market": "DE"}, thread1) - ] == [ - { - "__interrupt__": ( - Interrupt( - value={"value": 1}, - resumable=True, - ns=[AnyStr("node:")], - when="during", - ), - ) - } - ] - - assert [ - event - async for event in graph.astream( - Command(resume="answer 1", update={"my_key": "foofoo"}), - thread1, - stream_mode="updates", + assert [ + e async for e in graph.astream({"my_key": "DE", "market": "DE"}, thread1) + ] == [ + { + "__interrupt__": ( + Interrupt( + value={"value": 1}, + resumable=True, + ns=[AnyStr("node:")], + when="during", + ), ) - ] == [ - { - "__interrupt__": ( - Interrupt( - value={"value": 2}, - resumable=True, - ns=[AnyStr("node:")], - when="during", - ), - ) - } - ] + } + ] - assert [ - event - async for event in graph.astream( - Command(resume="answer 2"), thread1, stream_mode="updates" + assert [ + event + async for event in graph.astream( + Command(resume="answer 1", update={"my_key": "foofoo"}), + thread1, + stream_mode="updates", + ) + ] == [ + { + "__interrupt__": ( + Interrupt( + value={"value": 2}, + resumable=True, + ns=[AnyStr("node:")], + when="during", + ), ) - ] == [ - {"node": {"my_key": "answer 1 answer 2"}}, - ] + } + ] + + assert [ + event + async for event in graph.astream( + Command(resume="answer 2"), thread1, stream_mode="updates" + ) + ] == [ + {"node": {"my_key": "answer 1 answer 2"}}, + ] @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_interrupt_loop(checkpointer_name: str): +async def test_interrupt_loop(async_checkpointer: BaseCheckpointSaver) -> None: class State(TypedDict): age: int other: str @@ -6815,71 +6715,67 @@ async def test_interrupt_loop(checkpointer_name: str): builder.add_node("node", ask_age) builder.add_edge(START, "node") - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) - thread1 = {"configurable": {"thread_id": "1"}} + graph = builder.compile(checkpointer=async_checkpointer) + thread1 = {"configurable": {"thread_id": "1"}} - assert [e async for e in graph.astream({"other": ""}, thread1)] == [ - { - "__interrupt__": ( - Interrupt( - value="How old are you?", - resumable=True, - ns=[AnyStr("node:")], - when="during", - ), - ) - } - ] - - assert [ - event - async for event in graph.astream( - Command(resume="13"), - thread1, + assert [e async for e in graph.astream({"other": ""}, thread1)] == [ + { + "__interrupt__": ( + Interrupt( + value="How old are you?", + resumable=True, + ns=[AnyStr("node:")], + when="during", + ), ) - ] == [ - { - "__interrupt__": ( - Interrupt( - value="invalid response", - resumable=True, - ns=[AnyStr("node:")], - when="during", - ), - ) - } - ] + } + ] - assert [ - event - async for event in graph.astream( - Command(resume="15"), - thread1, + assert [ + event + async for event in graph.astream( + Command(resume="13"), + thread1, + ) + ] == [ + { + "__interrupt__": ( + Interrupt( + value="invalid response", + resumable=True, + ns=[AnyStr("node:")], + when="during", + ), ) - ] == [ - { - "__interrupt__": ( - Interrupt( - value="invalid response", - resumable=True, - ns=[AnyStr("node:")], - when="during", - ), - ) - } - ] + } + ] - assert [ - event async for event in graph.astream(Command(resume="19"), thread1) - ] == [ - {"node": {"age": 19}}, - ] + assert [ + event + async for event in graph.astream( + Command(resume="15"), + thread1, + ) + ] == [ + { + "__interrupt__": ( + Interrupt( + value="invalid response", + resumable=True, + ns=[AnyStr("node:")], + when="during", + ), + ) + } + ] + + assert [event async for event in graph.astream(Command(resume="19"), thread1)] == [ + {"node": {"age": 19}}, + ] @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_interrupt_functional(checkpointer_name: str) -> None: +async def test_interrupt_functional(async_checkpointer: BaseCheckpointSaver) -> None: @task async def foo(state: dict) -> dict: return {"a": state["a"] + "foo"} @@ -6888,27 +6784,26 @@ async def test_interrupt_functional(checkpointer_name: str) -> None: async def bar(state: dict) -> dict: return {"a": state["a"] + "bar", "b": state["b"]} - async with awith_checkpointer(checkpointer_name) as checkpointer: + @entrypoint(checkpointer=async_checkpointer) + async def graph(inputs: dict) -> dict: + foo_result = await foo(inputs) + value = interrupt("Provide value for bar:") + bar_input = {**foo_result, "b": value} + bar_result = await bar(bar_input) + return bar_result - @entrypoint(checkpointer=checkpointer) - async def graph(inputs: dict) -> dict: - foo_result = await foo(inputs) - value = interrupt("Provide value for bar:") - bar_input = {**foo_result, "b": value} - bar_result = await bar(bar_input) - return bar_result - - config = {"configurable": {"thread_id": "1"}} - # First run, interrupted at bar - await graph.ainvoke({"a": ""}, config) - # Resume with an answer - res = await graph.ainvoke(Command(resume="bar"), config) - assert res == {"a": "foobar", "b": "bar"} + config = {"configurable": {"thread_id": "1"}} + # First run, interrupted at bar + await graph.ainvoke({"a": ""}, config) + # Resume with an answer + res = await graph.ainvoke(Command(resume="bar"), config) + assert res == {"a": "foobar", "b": "bar"} @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_interrupt_task_functional(checkpointer_name: str) -> None: +async def test_interrupt_task_functional( + async_checkpointer: BaseCheckpointSaver, +) -> None: @task async def foo(state: dict) -> dict: return {"a": state["a"] + "foo"} @@ -6918,32 +6813,31 @@ async def test_interrupt_task_functional(checkpointer_name: str) -> None: value = interrupt("Provide value for bar:") return {"a": state["a"] + value} - async with awith_checkpointer(checkpointer_name) as checkpointer: + @entrypoint(checkpointer=async_checkpointer) + async def graph(inputs: dict) -> dict: + foo_result = await foo(inputs) + bar_result = await bar(foo_result) + return bar_result - @entrypoint(checkpointer=checkpointer) - async def graph(inputs: dict) -> dict: - foo_result = await foo(inputs) - bar_result = await bar(foo_result) - return bar_result - - config = {"configurable": {"thread_id": "1"}} - # First run, interrupted at bar - assert await graph.ainvoke({"a": ""}, config) == { - "__interrupt__": [ - Interrupt( - value="Provide value for bar:", - resumable=True, - ns=[AnyStr("graph:"), AnyStr("bar:")], - ), - ] - } - # Resume with an answer - res = await graph.ainvoke(Command(resume="bar"), config) - assert res == {"a": "foobar"} + config = {"configurable": {"thread_id": "1"}} + # First run, interrupted at bar + assert await graph.ainvoke({"a": ""}, config) == { + "__interrupt__": [ + Interrupt( + value="Provide value for bar:", + resumable=True, + ns=[AnyStr("graph:"), AnyStr("bar:")], + ), + ] + } + # Resume with an answer + res = await graph.ainvoke(Command(resume="bar"), config) + assert res == {"a": "foobar"} -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_command_with_static_breakpoints(checkpointer_name: str) -> None: +async def test_command_with_static_breakpoints( + async_checkpointer: BaseCheckpointSaver, +) -> None: """Test that we can use Command to resume and update with static breakpoints.""" class State(TypedDict): @@ -6967,18 +6861,16 @@ async def test_command_with_static_breakpoints(checkpointer_name: str) -> None: builder.add_edge(START, "node1") builder.add_edge("node1", "node2") - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["node1"]) - config = {"configurable": {"thread_id": str(uuid.uuid4())}} + graph = builder.compile(checkpointer=async_checkpointer, interrupt_before=["node1"]) + config = {"configurable": {"thread_id": str(uuid.uuid4())}} - # Start the graph and interrupt at the first node - await graph.ainvoke({"foo": "abc"}, config) - result = await graph.ainvoke(Command(update={"foo": "def"}), config) - assert result == {"foo": "def|node-1|node-2"} + # Start the graph and interrupt at the first node + await graph.ainvoke({"foo": "abc"}, config) + result = await graph.ainvoke(Command(update={"foo": "def"}), config) + assert result == {"foo": "def|node-1|node-2"} -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_multistep_plan(checkpointer_name: str): +async def test_multistep_plan(async_checkpointer: BaseCheckpointSaver) -> None: from langchain_core.messages import AnyMessage class State(TypedDict, total=False): @@ -7021,25 +6913,25 @@ async def test_multistep_plan(checkpointer_name: str): builder.add_node(step4) builder.add_edge(START, "planner") - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) + graph = builder.compile(checkpointer=async_checkpointer) - config = {"configurable": {"thread_id": "1"}} + config = {"configurable": {"thread_id": "1"}} - assert await graph.ainvoke({"messages": [("human", "start")]}, config) == { - "messages": [ - _AnyIdHumanMessage(content="start"), - _AnyIdHumanMessage(content="step1"), - _AnyIdHumanMessage(content="step2"), - _AnyIdHumanMessage(content="step3"), - _AnyIdHumanMessage(content="step4"), - ], - "plan": [], - } + assert await graph.ainvoke({"messages": [("human", "start")]}, config) == { + "messages": [ + _AnyIdHumanMessage(content="start"), + _AnyIdHumanMessage(content="step1"), + _AnyIdHumanMessage(content="step2"), + _AnyIdHumanMessage(content="step3"), + _AnyIdHumanMessage(content="step4"), + ], + "plan": [], + } -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_command_goto_with_static_breakpoints(checkpointer_name: str) -> None: +async def test_command_goto_with_static_breakpoints( + async_checkpointer: BaseCheckpointSaver, +) -> None: """Use Command goto with static breakpoints.""" class State(TypedDict): @@ -7063,15 +6955,14 @@ async def test_command_goto_with_static_breakpoints(checkpointer_name: str) -> N builder.add_edge(START, "node1") builder.add_edge("node1", "node2") - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer, interrupt_before=["node1"]) + graph = builder.compile(checkpointer=async_checkpointer, interrupt_before=["node1"]) - config = {"configurable": {"thread_id": str(uuid.uuid4())}} + config = {"configurable": {"thread_id": str(uuid.uuid4())}} - # Start the graph and interrupt at the first node - await graph.ainvoke({"foo": "abc"}, config) - result = await graph.ainvoke(Command(goto=["node2"]), config) - assert result == {"foo": "abc|node-1|node-2|node-2"} + # Start the graph and interrupt at the first node + await graph.ainvoke({"foo": "abc"}, config) + result = await graph.ainvoke(Command(goto=["node2"]), config) + assert result == {"foo": "abc|node-1|node-2|node-2"} async def test_parallel_node_execution(): @@ -7108,8 +6999,9 @@ async def test_parallel_node_execution(): @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_multiple_interrupt_state_persistence(checkpointer_name: str) -> None: +async def test_multiple_interrupt_state_persistence( + async_checkpointer: BaseCheckpointSaver, +) -> None: """Test that state is preserved correctly across multiple interrupts.""" class State(TypedDict): @@ -7124,31 +7016,30 @@ async def test_multiple_interrupt_state_persistence(checkpointer_name: str) -> N builder.add_node("node", interruptible_node) builder.add_edge(START, "node") - async with awith_checkpointer(checkpointer_name) as checkpointer: - app = builder.compile(checkpointer=checkpointer) - config = {"configurable": {"thread_id": "1"}} + app = builder.compile(checkpointer=async_checkpointer) + config = {"configurable": {"thread_id": "1"}} - # First execution - should hit first interrupt - await app.ainvoke({"steps": []}, config) + # First execution - should hit first interrupt + await app.ainvoke({"steps": []}, config) - # State should still be empty since node hasn't returned - state = await app.aget_state(config) - assert state.values == {"steps": []} + # State should still be empty since node hasn't returned + state = await app.aget_state(config) + assert state.values == {"steps": []} - # Resume after first interrupt - should hit second interrupt - await app.ainvoke(Command(resume="step1"), config) + # Resume after first interrupt - should hit second interrupt + await app.ainvoke(Command(resume="step1"), config) - # State should still be empty since node hasn't returned - state = await app.aget_state(config) - assert state.values == {"steps": []} + # State should still be empty since node hasn't returned + state = await app.aget_state(config) + assert state.values == {"steps": []} - # Resume after second interrupt - node should complete - result = await app.ainvoke(Command(resume="step2"), config) + # Resume after second interrupt - node should complete + result = await app.ainvoke(Command(resume="step2"), config) - # Now state should contain both steps since node returned - assert result["steps"] == ["step1", "step2"] - state = await app.aget_state(config) - assert state.values["steps"] == ["step1", "step2"] + # Now state should contain both steps since node returned + assert result["steps"] == ["step1", "step2"] + state = await app.aget_state(config) + assert state.values["steps"] == ["step1", "step2"] async def test_concurrent_execution(): @@ -7182,8 +7073,9 @@ async def test_concurrent_execution(): assert result["counter"] == 1 -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_checkpoint_recovery_async(checkpointer_name: str): +async def test_checkpoint_recovery_async( + async_checkpointer: BaseCheckpointSaver, +) -> None: """Test recovery from checkpoints after failures with async nodes.""" class State(TypedDict): @@ -7207,31 +7099,30 @@ async def test_checkpoint_recovery_async(checkpointer_name: str): builder.add_edge(START, "node1") builder.add_edge("node1", "node2") - async with awith_checkpointer(checkpointer_name) as checkpointer: - graph = builder.compile(checkpointer=checkpointer) - config = {"configurable": {"thread_id": "1"}} + graph = builder.compile(checkpointer=async_checkpointer) + config = {"configurable": {"thread_id": "1"}} - # First attempt should fail - with pytest.raises(RuntimeError): - await graph.ainvoke({"steps": ["start"], "attempt": 1}, config) + # First attempt should fail + with pytest.raises(RuntimeError): + await graph.ainvoke({"steps": ["start"], "attempt": 1}, config) - # Verify checkpoint state - state = await graph.aget_state(config) - assert state is not None - assert state.values == {"steps": ["start"], "attempt": 1} # input state saved - assert state.next == ("node1",) # Should retry failed node + # Verify checkpoint state + state = await graph.aget_state(config) + assert state is not None + assert state.values == {"steps": ["start"], "attempt": 1} # input state saved + assert state.next == ("node1",) # Should retry failed node - # Retry with updated attempt count - result = await graph.ainvoke({"steps": [], "attempt": 2}, config) - assert result == {"steps": ["start", "node1", "node2"], "attempt": 2} + # Retry with updated attempt count + result = await graph.ainvoke({"steps": [], "attempt": 2}, config) + assert result == {"steps": ["start", "node1", "node2"], "attempt": 2} - # Verify checkpoint history shows both attempts - history = [c async for c in graph.aget_state_history(config)] - assert len(history) == 6 # Initial + failed attempt + successful attempt + # Verify checkpoint history shows both attempts + history = [c async for c in graph.aget_state_history(config)] + assert len(history) == 6 # Initial + failed attempt + successful attempt - # Verify the error was recorded in checkpoint - failed_checkpoint = next(c for c in history if c.tasks and c.tasks[0].error) - assert "RuntimeError('Simulated failure')" in failed_checkpoint.tasks[0].error + # Verify the error was recorded in checkpoint + failed_checkpoint = next(c for c in history if c.tasks and c.tasks[0].error) + assert "RuntimeError('Simulated failure')" in failed_checkpoint.tasks[0].error async def test_multiple_updates_root() -> None: @@ -7286,30 +7177,28 @@ async def test_multiple_updates() -> None: @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_falsy_return_from_task(checkpointer_name: str) -> None: +async def test_falsy_return_from_task(async_checkpointer: BaseCheckpointSaver) -> None: """Test with a falsy return from a task.""" @task async def falsy_task() -> bool: return False - async with awith_checkpointer(checkpointer_name) as checkpointer: + @entrypoint(checkpointer=async_checkpointer) + async def graph(state: dict) -> dict: + """React tool.""" + await falsy_task() + interrupt("test") - @entrypoint(checkpointer=checkpointer) - async def graph(state: dict) -> dict: - """React tool.""" - await falsy_task() - interrupt("test") - - configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} - await graph.ainvoke({"a": 5}, configurable) - await graph.ainvoke(Command(resume="123"), configurable) + configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} + await graph.ainvoke({"a": 5}, configurable) + await graph.ainvoke(Command(resume="123"), configurable) @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_multiple_interrupts_functional(checkpointer_name: str) -> None: +async def test_multiple_interrupts_functional( + async_checkpointer: BaseCheckpointSaver, +) -> None: """Test multiple interrupts with functional API.""" from langgraph.func import entrypoint, task @@ -7322,108 +7211,105 @@ async def test_multiple_interrupts_functional(checkpointer_name: str) -> None: counter += 1 return 2 * x - async with awith_checkpointer(checkpointer_name) as checkpointer: + @entrypoint(checkpointer=async_checkpointer) + async def graph(state: dict) -> dict: + """React tool.""" - @entrypoint(checkpointer=checkpointer) - async def graph(state: dict) -> dict: - """React tool.""" + values = [] - values = [] + for idx in [1, 2, 3]: + values.extend([await double(idx), interrupt({"a": "boo"})]) - for idx in [1, 2, 3]: - values.extend([await double(idx), interrupt({"a": "boo"})]) + return {"values": values} - 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 + 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 @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_multiple_interrupts_functional_cache( - checkpointer_name: str, cache: BaseCache + async_checkpointer: BaseCheckpointSaver, cache: BaseCache ): """Test multiple interrupts with functional API.""" - async with awith_checkpointer(checkpointer_name) as checkpointer: - counter = 0 + counter = 0 - @task(cache_policy=CachePolicy()) - def double(x: int) -> int: - """Increment the counter.""" - nonlocal counter - counter += 1 - return 2 * x + @task(cache_policy=CachePolicy()) + def double(x: int) -> int: + """Increment the counter.""" + nonlocal counter + counter += 1 + return 2 * x - @entrypoint(checkpointer=checkpointer, cache=cache) - def graph(state: dict) -> dict: - """React tool.""" + @entrypoint(checkpointer=async_checkpointer, cache=cache) + def graph(state: dict) -> dict: + """React tool.""" - values = [] + values = [] - for idx in [1, 1, 2, 2, 3, 3]: - values.extend([double(idx).result(), interrupt({"a": "boo"})]) + for idx in [1, 1, 2, 2, 3, 3]: + values.extend([double(idx).result(), interrupt({"a": "boo"})]) - return {"values": values} + 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) - await graph.ainvoke(Command(resume="c"), configurable) - await graph.ainvoke(Command(resume="d"), configurable) - await graph.ainvoke(Command(resume="e"), configurable) - result = await graph.ainvoke(Command(resume="f"), configurable) - # `double` value should be cached appropriately when used w/ `interrupt` - assert result == { - "values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"], - } - assert counter == 3 + 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) + await graph.ainvoke(Command(resume="c"), configurable) + await graph.ainvoke(Command(resume="d"), configurable) + await graph.ainvoke(Command(resume="e"), configurable) + result = await graph.ainvoke(Command(resume="f"), configurable) + # `double` value should be cached appropriately when used w/ `interrupt` + assert result == { + "values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"], + } + assert counter == 3 - 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) - await graph.ainvoke(Command(resume="c"), configurable) - await graph.ainvoke(Command(resume="d"), configurable) - await graph.ainvoke(Command(resume="e"), configurable) - result = await graph.ainvoke(Command(resume="f"), configurable) - # `double` value should be cached appropriately when used w/ `interrupt` - assert result == { - "values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"], - } - assert counter == 3 + 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) + await graph.ainvoke(Command(resume="c"), configurable) + await graph.ainvoke(Command(resume="d"), configurable) + await graph.ainvoke(Command(resume="e"), configurable) + result = await graph.ainvoke(Command(resume="f"), configurable) + # `double` value should be cached appropriately when used w/ `interrupt` + assert result == { + "values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"], + } + assert counter == 3 - # clear the cache - await double.aclear_cache(cache) + # clear the cache + await double.aclear_cache(cache) - # now should recompute - 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) - await graph.ainvoke(Command(resume="c"), configurable) - await graph.ainvoke(Command(resume="d"), configurable) - await graph.ainvoke(Command(resume="e"), configurable) - result = await graph.ainvoke(Command(resume="f"), configurable) - # `double` value should be cached appropriately when used w/ `interrupt` - assert result == { - "values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"], - } - assert counter == 6 + # now should recompute + 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) + await graph.ainvoke(Command(resume="c"), configurable) + await graph.ainvoke(Command(resume="d"), configurable) + await graph.ainvoke(Command(resume="e"), configurable) + result = await graph.ainvoke(Command(resume="f"), configurable) + # `double` value should be cached appropriately when used w/ `interrupt` + assert result == { + "values": [2, "a", 2, "b", 4, "c", 4, "d", 6, "e", 6, "f"], + } + assert counter == 6 @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_double_interrupt_subgraph(checkpointer_name: str) -> None: +async def test_double_interrupt_subgraph( + async_checkpointer: BaseCheckpointSaver, +) -> None: class AgentState(TypedDict): input: str @@ -7444,97 +7330,92 @@ async def test_double_interrupt_subgraph(checkpointer_name: str) -> None: .add_edge("node_2", END) ) - async with awith_checkpointer(checkpointer_name) as checkpointer: - # invoke the sub graph - subgraph = subgraph_builder.compile(checkpointer=checkpointer) - thread = {"configurable": {"thread_id": str(uuid.uuid4())}} - assert [c async for c in subgraph.astream({"input": "test"}, thread)] == [ - { - "__interrupt__": ( - Interrupt( - value="interrupt node 1", - resumable=True, - ns=[AnyStr("node_1:")], - when="during", - ), - ) - }, - ] - # resume from the first interrupt - assert [c async for c in subgraph.astream(Command(resume="123"), thread)] == [ - { - "node_1": {"input": "123"}, - }, - { - "__interrupt__": ( - Interrupt( - value="interrupt node 2", - resumable=True, - ns=[AnyStr("node_2:")], - when="during", - ), - ) - }, - ] - # resume from the second interrupt - assert [c async for c in subgraph.astream(Command(resume="123"), thread)] == [ - { - "node_2": {"input": "123"}, - }, - ] + # invoke the sub graph + subgraph = subgraph_builder.compile(checkpointer=async_checkpointer) + thread = {"configurable": {"thread_id": str(uuid.uuid4())}} + assert [c async for c in subgraph.astream({"input": "test"}, thread)] == [ + { + "__interrupt__": ( + Interrupt( + value="interrupt node 1", + resumable=True, + ns=[AnyStr("node_1:")], + when="during", + ), + ) + }, + ] + # resume from the first interrupt + assert [c async for c in subgraph.astream(Command(resume="123"), thread)] == [ + { + "node_1": {"input": "123"}, + }, + { + "__interrupt__": ( + Interrupt( + value="interrupt node 2", + resumable=True, + ns=[AnyStr("node_2:")], + when="during", + ), + ) + }, + ] + # resume from the second interrupt + assert [c async for c in subgraph.astream(Command(resume="123"), thread)] == [ + { + "node_2": {"input": "123"}, + }, + ] - subgraph = subgraph_builder.compile() + subgraph = subgraph_builder.compile() - def invoke_sub_agent(state: AgentState): - return subgraph.invoke(state) + def invoke_sub_agent(state: AgentState): + return subgraph.invoke(state) - thread = {"configurable": {"thread_id": str(uuid.uuid4())}} + thread = {"configurable": {"thread_id": str(uuid.uuid4())}} - parent_agent = ( - StateGraph(AgentState) - .add_node("invoke_sub_agent", invoke_sub_agent) - .add_edge(START, "invoke_sub_agent") - .add_edge("invoke_sub_agent", END) - .compile(checkpointer=checkpointer) - ) + parent_agent = ( + StateGraph(AgentState) + .add_node("invoke_sub_agent", invoke_sub_agent) + .add_edge(START, "invoke_sub_agent") + .add_edge("invoke_sub_agent", END) + .compile(checkpointer=async_checkpointer) + ) - assert [c async for c in parent_agent.astream({"input": "test"}, thread)] == [ - { - "__interrupt__": ( - Interrupt( - value="interrupt node 1", - resumable=True, - ns=[AnyStr("invoke_sub_agent:"), AnyStr("node_1:")], - when="during", - ), - ) - }, - ] + assert [c async for c in parent_agent.astream({"input": "test"}, thread)] == [ + { + "__interrupt__": ( + Interrupt( + value="interrupt node 1", + resumable=True, + ns=[AnyStr("invoke_sub_agent:"), AnyStr("node_1:")], + when="during", + ), + ) + }, + ] - # resume from the first interrupt - assert [ - c async for c in parent_agent.astream(Command(resume=True), thread) - ] == [ - { - "__interrupt__": ( - Interrupt( - value="interrupt node 2", - resumable=True, - ns=[AnyStr("invoke_sub_agent:"), AnyStr("node_2:")], - when="during", - ), - ) - } - ] + # resume from the first interrupt + assert [c async for c in parent_agent.astream(Command(resume=True), thread)] == [ + { + "__interrupt__": ( + Interrupt( + value="interrupt node 2", + resumable=True, + ns=[AnyStr("invoke_sub_agent:"), AnyStr("node_2:")], + when="during", + ), + ) + } + ] - # resume from 2nd interrupt - assert [ - c async for c in parent_agent.astream(Command(resume=True), thread) - ] == [ - { - "invoke_sub_agent": {"input": True}, - }, - ] + # resume from 2nd interrupt + assert [c async for c in parent_agent.astream(Command(resume=True), thread)] == [ + { + "invoke_sub_agent": {"input": True}, + }, + ] @NEEDS_CONTEXTVARS @@ -7575,8 +7456,7 @@ async def test_async_streaming_with_functional_api() -> None: @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_multiple_subgraphs(checkpointer_name: str) -> None: +async def test_multiple_subgraphs(async_checkpointer: BaseCheckpointSaver) -> None: class State(TypedDict): a: int b: int @@ -7584,126 +7464,119 @@ async def test_multiple_subgraphs(checkpointer_name: str) -> None: class Output(TypedDict): result: int - async with awith_checkpointer(checkpointer_name) as checkpointer: - # Define the subgraphs - async def add(state): - return {"result": state["a"] + state["b"]} + # Define the subgraphs + async def add(state): + return {"result": state["a"] + state["b"]} - add_subgraph = ( - StateGraph(State, output=Output) - .add_node(add) - .add_edge(START, "add") - .compile() - ) + add_subgraph = ( + StateGraph(State, output=Output).add_node(add).add_edge(START, "add").compile() + ) - async def multiply(state): - return {"result": state["a"] * state["b"]} + async def multiply(state): + return {"result": state["a"] * state["b"]} - multiply_subgraph = ( - StateGraph(State, output=Output) - .add_node(multiply) - .add_edge(START, "multiply") - .compile() - ) + multiply_subgraph = ( + StateGraph(State, output=Output) + .add_node(multiply) + .add_edge(START, "multiply") + .compile() + ) - # Test calling the same subgraph multiple times - async def call_same_subgraph(state): - result = await add_subgraph.ainvoke(state) - another_result = await add_subgraph.ainvoke( - {"a": result["result"], "b": 10} - ) - return another_result + # Test calling the same subgraph multiple times + async def call_same_subgraph(state): + result = await add_subgraph.ainvoke(state) + another_result = await add_subgraph.ainvoke({"a": result["result"], "b": 10}) + return another_result - parent_call_same_subgraph = ( - StateGraph(State, output=Output) - .add_node(call_same_subgraph) - .add_edge(START, "call_same_subgraph") - .compile(checkpointer=checkpointer) - ) - config = {"configurable": {"thread_id": "1"}} - assert await parent_call_same_subgraph.ainvoke({"a": 2, "b": 3}, config) == { - "result": 15 + parent_call_same_subgraph = ( + StateGraph(State, output=Output) + .add_node(call_same_subgraph) + .add_edge(START, "call_same_subgraph") + .compile(checkpointer=async_checkpointer) + ) + config = {"configurable": {"thread_id": "1"}} + assert await parent_call_same_subgraph.ainvoke({"a": 2, "b": 3}, config) == { + "result": 15 + } + + # Test calling multiple subgraphs + class Output(TypedDict): + add_result: int + multiply_result: int + + async def call_multiple_subgraphs(state): + add_result = await add_subgraph.ainvoke(state) + multiply_result = await multiply_subgraph.ainvoke(state) + return { + "add_result": add_result["result"], + "multiply_result": multiply_result["result"], } - # Test calling multiple subgraphs - class Output(TypedDict): - add_result: int - multiply_result: int - - async def call_multiple_subgraphs(state): - add_result = await add_subgraph.ainvoke(state) - multiply_result = await multiply_subgraph.ainvoke(state) - return { - "add_result": add_result["result"], - "multiply_result": multiply_result["result"], - } - - parent_call_multiple_subgraphs = ( - StateGraph(State, output=Output) - .add_node(call_multiple_subgraphs) - .add_edge(START, "call_multiple_subgraphs") - .compile(checkpointer=checkpointer) - ) - config = {"configurable": {"thread_id": "2"}} - assert await parent_call_multiple_subgraphs.ainvoke( - {"a": 2, "b": 3}, config - ) == { - "add_result": 5, - "multiply_result": 6, - } + parent_call_multiple_subgraphs = ( + StateGraph(State, output=Output) + .add_node(call_multiple_subgraphs) + .add_edge(START, "call_multiple_subgraphs") + .compile(checkpointer=async_checkpointer) + ) + config = {"configurable": {"thread_id": "2"}} + assert await parent_call_multiple_subgraphs.ainvoke({"a": 2, "b": 3}, config) == { + "add_result": 5, + "multiply_result": 6, + } @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_multiple_subgraphs_functional(checkpointer_name: str) -> None: - async with awith_checkpointer(checkpointer_name) as checkpointer: - # Define addition subgraph - @entrypoint() - async def add(inputs): - a, b = inputs - return a + b +async def test_multiple_subgraphs_functional( + async_checkpointer: BaseCheckpointSaver, +) -> None: + # Define addition subgraph + @entrypoint() + async def add(inputs): + a, b = inputs + return a + b - # Define multiplication subgraph using tasks - @task - async def multiply_task(a, b): - return a * b + # Define multiplication subgraph using tasks + @task + async def multiply_task(a, b): + return a * b - @entrypoint() - async def multiply(inputs): - return await multiply_task(*inputs) + @entrypoint() + async def multiply(inputs): + return await multiply_task(*inputs) - # Test calling the same subgraph multiple times - @task - async def call_same_subgraph(a, b): - result = await add.ainvoke([a, b]) - another_result = await add.ainvoke([result, 10]) - return another_result + # Test calling the same subgraph multiple times + @task + async def call_same_subgraph(a, b): + result = await add.ainvoke([a, b]) + another_result = await add.ainvoke([result, 10]) + return another_result - @entrypoint(checkpointer=checkpointer) - async def parent_call_same_subgraph(inputs): - return await call_same_subgraph(*inputs) + @entrypoint(checkpointer=async_checkpointer) + async def parent_call_same_subgraph(inputs): + return await call_same_subgraph(*inputs) - config = {"configurable": {"thread_id": "1"}} - assert await parent_call_same_subgraph.ainvoke([2, 3], config) == 15 + config = {"configurable": {"thread_id": "1"}} + assert await parent_call_same_subgraph.ainvoke([2, 3], config) == 15 - # Test calling multiple subgraphs - @task - async def call_multiple_subgraphs(a, b): - add_result = await add.ainvoke([a, b]) - multiply_result = await multiply.ainvoke([a, b]) - return [add_result, multiply_result] + # Test calling multiple subgraphs + @task + async def call_multiple_subgraphs(a, b): + add_result = await add.ainvoke([a, b]) + multiply_result = await multiply.ainvoke([a, b]) + return [add_result, multiply_result] - @entrypoint(checkpointer=checkpointer) - async def parent_call_multiple_subgraphs(inputs): - return await call_multiple_subgraphs(*inputs) + @entrypoint(checkpointer=async_checkpointer) + async def parent_call_multiple_subgraphs(inputs): + return await call_multiple_subgraphs(*inputs) - config = {"configurable": {"thread_id": "2"}} - assert await parent_call_multiple_subgraphs.ainvoke([2, 3], config) == [5, 6] + config = {"configurable": {"thread_id": "2"}} + assert await parent_call_multiple_subgraphs.ainvoke([2, 3], config) == [5, 6] @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_multiple_subgraphs_mixed_entrypoint(checkpointer_name: str) -> None: +async def test_multiple_subgraphs_mixed_entrypoint( + async_checkpointer: BaseCheckpointSaver, +) -> None: """Test calling multiple StateGraph subgraphs from an entrypoint.""" class State(TypedDict): @@ -7713,218 +7586,205 @@ async def test_multiple_subgraphs_mixed_entrypoint(checkpointer_name: str) -> No class Output(TypedDict): result: int - async with awith_checkpointer(checkpointer_name) as checkpointer: - # Define the subgraphs - async def add(state): - return {"result": state["a"] + state["b"]} + # Define the subgraphs + async def add(state): + return {"result": state["a"] + state["b"]} - add_subgraph = ( - StateGraph(State, output=Output) - .add_node(add) - .add_edge(START, "add") - .compile() - ) + add_subgraph = ( + StateGraph(State, output=Output).add_node(add).add_edge(START, "add").compile() + ) - async def multiply(state): - return {"result": state["a"] * state["b"]} + async def multiply(state): + return {"result": state["a"] * state["b"]} - multiply_subgraph = ( - StateGraph(State, output=Output) - .add_node(multiply) - .add_edge(START, "multiply") - .compile() - ) + multiply_subgraph = ( + StateGraph(State, output=Output) + .add_node(multiply) + .add_edge(START, "multiply") + .compile() + ) - # Test calling the same subgraph multiple times - @task - async def call_same_subgraph(a, b): - result = (await add_subgraph.ainvoke({"a": a, "b": b}))["result"] - another_result = (await add_subgraph.ainvoke({"a": result, "b": 10}))[ - "result" - ] - return another_result + # Test calling the same subgraph multiple times + @task + async def call_same_subgraph(a, b): + result = (await add_subgraph.ainvoke({"a": a, "b": b}))["result"] + another_result = (await add_subgraph.ainvoke({"a": result, "b": 10}))["result"] + return another_result - @entrypoint(checkpointer=checkpointer) - async def parent_call_same_subgraph(inputs): - return await call_same_subgraph(*inputs) + @entrypoint(checkpointer=async_checkpointer) + async def parent_call_same_subgraph(inputs): + return await call_same_subgraph(*inputs) - config = {"configurable": {"thread_id": "1"}} - assert await parent_call_same_subgraph.ainvoke([2, 3], config) == 15 + config = {"configurable": {"thread_id": "1"}} + assert await parent_call_same_subgraph.ainvoke([2, 3], config) == 15 - # Test calling multiple subgraphs - @task - async def call_multiple_subgraphs(a, b): - add_result = (await add_subgraph.ainvoke({"a": a, "b": b}))["result"] - multiply_result = (await multiply_subgraph.ainvoke({"a": a, "b": b}))[ - "result" - ] - return [add_result, multiply_result] + # Test calling multiple subgraphs + @task + async def call_multiple_subgraphs(a, b): + add_result = (await add_subgraph.ainvoke({"a": a, "b": b}))["result"] + multiply_result = (await multiply_subgraph.ainvoke({"a": a, "b": b}))["result"] + return [add_result, multiply_result] - @entrypoint(checkpointer=checkpointer) - async def parent_call_multiple_subgraphs(inputs): - return await call_multiple_subgraphs(*inputs) + @entrypoint(checkpointer=async_checkpointer) + async def parent_call_multiple_subgraphs(inputs): + return await call_multiple_subgraphs(*inputs) - config = {"configurable": {"thread_id": "2"}} - assert await parent_call_multiple_subgraphs.ainvoke([2, 3], config) == [5, 6] + config = {"configurable": {"thread_id": "2"}} + assert await parent_call_multiple_subgraphs.ainvoke([2, 3], config) == [5, 6] @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_multiple_subgraphs_mixed_state_graph( - request: pytest.FixtureRequest, checkpointer_name: str + async_checkpointer: BaseCheckpointSaver, ) -> None: """Test calling multiple entrypoint "subgraphs" from a StateGraph.""" - async with awith_checkpointer(checkpointer_name) as checkpointer: - class State(TypedDict): - a: int - b: int + class State(TypedDict): + a: int + b: int - class Output(TypedDict): - result: int + class Output(TypedDict): + result: int - # Define addition subgraph - @entrypoint() - async def add(inputs): - a, b = inputs - return a + b + # Define addition subgraph + @entrypoint() + async def add(inputs): + a, b = inputs + return a + b - # Define multiplication subgraph using tasks - @task - async def multiply_task(a, b): - return a * b + # Define multiplication subgraph using tasks + @task + async def multiply_task(a, b): + return a * b - @entrypoint() - async def multiply(inputs): - return await multiply_task(*inputs) + @entrypoint() + async def multiply(inputs): + return await multiply_task(*inputs) - # Test calling the same subgraph multiple times - async def call_same_subgraph(state): - result = await add.ainvoke([state["a"], state["b"]]) - another_result = await add.ainvoke([result, 10]) - return {"result": another_result} + # Test calling the same subgraph multiple times + async def call_same_subgraph(state): + result = await add.ainvoke([state["a"], state["b"]]) + another_result = await add.ainvoke([result, 10]) + return {"result": another_result} - parent_call_same_subgraph = ( - StateGraph(State, output=Output) - .add_node(call_same_subgraph) - .add_edge(START, "call_same_subgraph") - .compile(checkpointer=checkpointer) - ) - config = {"configurable": {"thread_id": "1"}} - assert await parent_call_same_subgraph.ainvoke({"a": 2, "b": 3}, config) == { - "result": 15 + parent_call_same_subgraph = ( + StateGraph(State, output=Output) + .add_node(call_same_subgraph) + .add_edge(START, "call_same_subgraph") + .compile(checkpointer=async_checkpointer) + ) + config = {"configurable": {"thread_id": "1"}} + assert await parent_call_same_subgraph.ainvoke({"a": 2, "b": 3}, config) == { + "result": 15 + } + + # Test calling multiple subgraphs + class Output(TypedDict): + add_result: int + multiply_result: int + + async def call_multiple_subgraphs(state): + add_result = await add.ainvoke([state["a"], state["b"]]) + multiply_result = await multiply.ainvoke([state["a"], state["b"]]) + return { + "add_result": add_result, + "multiply_result": multiply_result, } - # Test calling multiple subgraphs - class Output(TypedDict): - add_result: int - multiply_result: int - - async def call_multiple_subgraphs(state): - add_result = await add.ainvoke([state["a"], state["b"]]) - multiply_result = await multiply.ainvoke([state["a"], state["b"]]) - return { - "add_result": add_result, - "multiply_result": multiply_result, - } - - parent_call_multiple_subgraphs = ( - StateGraph(State, output=Output) - .add_node(call_multiple_subgraphs) - .add_edge(START, "call_multiple_subgraphs") - .compile(checkpointer=checkpointer) - ) - config = {"configurable": {"thread_id": "2"}} - assert await parent_call_multiple_subgraphs.ainvoke( - {"a": 2, "b": 3}, config - ) == { - "add_result": 5, - "multiply_result": 6, - } + parent_call_multiple_subgraphs = ( + StateGraph(State, output=Output) + .add_node(call_multiple_subgraphs) + .add_edge(START, "call_multiple_subgraphs") + .compile(checkpointer=async_checkpointer) + ) + config = {"configurable": {"thread_id": "2"}} + assert await parent_call_multiple_subgraphs.ainvoke({"a": 2, "b": 3}, config) == { + "add_result": 5, + "multiply_result": 6, + } @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_multiple_subgraphs_checkpointer(checkpointer_name: str) -> None: - async with awith_checkpointer(checkpointer_name) as checkpointer: +async def test_multiple_subgraphs_checkpointer( + async_checkpointer: BaseCheckpointSaver, +) -> None: + class SubgraphState(TypedDict): + sub_counter: Annotated[int, operator.add] - class SubgraphState(TypedDict): - sub_counter: Annotated[int, operator.add] + async def subgraph_node(state): + return {"sub_counter": 2} - async def subgraph_node(state): - return {"sub_counter": 2} + sub_graph_1 = ( + StateGraph(SubgraphState) + .add_node(subgraph_node) + .add_edge(START, "subgraph_node") + .compile(checkpointer=True) + ) - sub_graph_1 = ( - StateGraph(SubgraphState) - .add_node(subgraph_node) - .add_edge(START, "subgraph_node") - .compile(checkpointer=True) + class OtherSubgraphState(TypedDict): + other_sub_counter: Annotated[int, operator.add] + + async def other_subgraph_node(state): + return {"other_sub_counter": 3} + + sub_graph_2 = ( + StateGraph(OtherSubgraphState) + .add_node(other_subgraph_node) + .add_edge(START, "other_subgraph_node") + .compile() + ) + + class ParentState(TypedDict): + parent_counter: int + + async def parent_node(state): + result = await sub_graph_1.ainvoke({"sub_counter": state["parent_counter"]}) + other_result = await sub_graph_2.ainvoke( + {"other_sub_counter": result["sub_counter"]} ) + return {"parent_counter": other_result["other_sub_counter"]} - class OtherSubgraphState(TypedDict): - other_sub_counter: Annotated[int, operator.add] + parent_graph = ( + StateGraph(ParentState) + .add_node(parent_node) + .add_edge(START, "parent_node") + .compile(checkpointer=async_checkpointer) + ) - async def other_subgraph_node(state): - return {"other_sub_counter": 3} - - sub_graph_2 = ( - StateGraph(OtherSubgraphState) - .add_node(other_subgraph_node) - .add_edge(START, "other_subgraph_node") - .compile() + config = {"configurable": {"thread_id": "1"}} + assert await parent_graph.ainvoke({"parent_counter": 0}, config) == { + "parent_counter": 5 + } + assert await parent_graph.ainvoke({"parent_counter": 0}, config) == { + "parent_counter": 7 + } + config = {"configurable": {"thread_id": "2"}} + assert [ + c + async for c in parent_graph.astream( + {"parent_counter": 0}, config, subgraphs=True, stream_mode="updates" ) - - class ParentState(TypedDict): - parent_counter: int - - async def parent_node(state): - result = await sub_graph_1.ainvoke({"sub_counter": state["parent_counter"]}) - other_result = await sub_graph_2.ainvoke( - {"other_sub_counter": result["sub_counter"]} - ) - return {"parent_counter": other_result["other_sub_counter"]} - - parent_graph = ( - StateGraph(ParentState) - .add_node(parent_node) - .add_edge(START, "parent_node") - .compile(checkpointer=checkpointer) + ] == [ + (("parent_node",), {"subgraph_node": {"sub_counter": 2}}), + ( + (AnyStr("parent_node:"), "1"), + {"other_subgraph_node": {"other_sub_counter": 3}}, + ), + ((), {"parent_node": {"parent_counter": 5}}), + ] + assert [ + c + async for c in parent_graph.astream( + {"parent_counter": 0}, config, subgraphs=True, stream_mode="updates" ) - - config = {"configurable": {"thread_id": "1"}} - assert await parent_graph.ainvoke({"parent_counter": 0}, config) == { - "parent_counter": 5 - } - assert await parent_graph.ainvoke({"parent_counter": 0}, config) == { - "parent_counter": 7 - } - config = {"configurable": {"thread_id": "2"}} - assert [ - c - async for c in parent_graph.astream( - {"parent_counter": 0}, config, subgraphs=True, stream_mode="updates" - ) - ] == [ - (("parent_node",), {"subgraph_node": {"sub_counter": 2}}), - ( - (AnyStr("parent_node:"), "1"), - {"other_subgraph_node": {"other_sub_counter": 3}}, - ), - ((), {"parent_node": {"parent_counter": 5}}), - ] - assert [ - c - async for c in parent_graph.astream( - {"parent_counter": 0}, config, subgraphs=True, stream_mode="updates" - ) - ] == [ - (("parent_node",), {"subgraph_node": {"sub_counter": 2}}), - ( - (AnyStr("parent_node:"), "1"), - {"other_subgraph_node": {"other_sub_counter": 3}}, - ), - ((), {"parent_node": {"parent_counter": 7}}), - ] + ] == [ + (("parent_node",), {"subgraph_node": {"sub_counter": 2}}), + ( + (AnyStr("parent_node:"), "1"), + {"other_subgraph_node": {"other_sub_counter": 3}}, + ), + ((), {"parent_node": {"parent_counter": 7}}), + ] @NEEDS_CONTEXTVARS @@ -8279,180 +8139,177 @@ async def test_stream_messages_dedupe_inputs() -> None: assert chunks[0][1]["langgraph_node"] == "call_model" -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_stream_messages_dedupe_state(checkpointer_name: str) -> None: - async with awith_checkpointer(checkpointer_name) as checkpointer: - from langchain_core.messages import AIMessage +async def test_stream_messages_dedupe_state( + async_checkpointer: BaseCheckpointSaver, +) -> None: + from langchain_core.messages import AIMessage - to_emit = [AIMessage("bye", id="1"), AIMessage("bye again", id="2")] + to_emit = [AIMessage("bye", id="1"), AIMessage("bye again", id="2")] - async def call_model(state): - return {"messages": to_emit.pop(0)} + async def call_model(state): + return {"messages": to_emit.pop(0)} - async def route(state): - return Command(goto="node_2", graph=Command.PARENT) + async def route(state): + return Command(goto="node_2", graph=Command.PARENT) - subgraph = ( - StateGraph(MessagesState) - .add_node(call_model) - .add_node(route) - .add_edge(START, "call_model") - .add_edge("call_model", "route") - .compile() + subgraph = ( + StateGraph(MessagesState) + .add_node(call_model) + .add_node(route) + .add_edge(START, "call_model") + .add_edge("call_model", "route") + .compile() + ) + + graph = ( + StateGraph(MessagesState) + .add_node("node_1", subgraph) + .add_node("node_2", lambda state: state) + .add_edge(START, "node_1") + .compile(checkpointer=async_checkpointer) + ) + + thread1 = {"configurable": {"thread_id": "1"}} + + chunks = [ + chunk + async for ns, chunk in graph.astream( + {"messages": "hi"}, thread1, stream_mode="messages", subgraphs=True ) + ] - graph = ( - StateGraph(MessagesState) - .add_node("node_1", subgraph) - .add_node("node_2", lambda state: state) - .add_edge(START, "node_1") - .compile(checkpointer=checkpointer) + assert len(chunks) == 1 + assert chunks[0][0] == AIMessage("bye", id="1") + assert chunks[0][1]["langgraph_node"] == "call_model" + + chunks = [ + chunk + async for ns, chunk in graph.astream( + {"messages": "hi again"}, + thread1, + stream_mode="messages", + subgraphs=True, ) + ] - thread1 = {"configurable": {"thread_id": "1"}} - - chunks = [ - chunk - async for ns, chunk in graph.astream( - {"messages": "hi"}, thread1, stream_mode="messages", subgraphs=True - ) - ] - - assert len(chunks) == 1 - assert chunks[0][0] == AIMessage("bye", id="1") - assert chunks[0][1]["langgraph_node"] == "call_model" - - chunks = [ - chunk - async for ns, chunk in graph.astream( - {"messages": "hi again"}, - thread1, - stream_mode="messages", - subgraphs=True, - ) - ] - - assert len(chunks) == 1 - assert chunks[0][0] == AIMessage("bye again", id="2") - assert chunks[0][1]["langgraph_node"] == "call_model" + assert len(chunks) == 1 + assert chunks[0][0] == AIMessage("bye again", id="2") + assert chunks[0][1]["langgraph_node"] == "call_model" @NEEDS_CONTEXTVARS -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_interrupt_subgraph_reenter_checkpointer_true( - checkpointer_name: str, + async_checkpointer: BaseCheckpointSaver, ) -> None: - async with awith_checkpointer(checkpointer_name) as checkpointer: + class SubgraphState(TypedDict): + foo: str + bar: str - class SubgraphState(TypedDict): - foo: str - bar: str + class ParentState(TypedDict): + foo: str + counter: int - class ParentState(TypedDict): - foo: str - counter: int + called = [] + bar_values = [] - called = [] - bar_values = [] + async def subnode_1(state: SubgraphState): + called.append("subnode_1") + bar_values.append(state.get("bar")) + return {"foo": "subgraph_1"} - async def subnode_1(state: SubgraphState): - called.append("subnode_1") - bar_values.append(state.get("bar")) - return {"foo": "subgraph_1"} + async def subnode_2(state: SubgraphState): + called.append("subnode_2") + value = interrupt("Provide value") + value += "baz" + return {"foo": "subgraph_2", "bar": value} - async def subnode_2(state: SubgraphState): - called.append("subnode_2") - value = interrupt("Provide value") - value += "baz" - return {"foo": "subgraph_2", "bar": value} + subgraph = ( + StateGraph(SubgraphState) + .add_node(subnode_1) + .add_node(subnode_2) + .add_edge(START, "subnode_1") + .add_edge("subnode_1", "subnode_2") + .compile(checkpointer=True) + ) - subgraph = ( - StateGraph(SubgraphState) - .add_node(subnode_1) - .add_node(subnode_2) - .add_edge(START, "subnode_1") - .add_edge("subnode_1", "subnode_2") - .compile(checkpointer=True) - ) + async def call_subgraph(state: ParentState): + called.append("call_subgraph") + return await subgraph.ainvoke(state) - async def call_subgraph(state: ParentState): - called.append("call_subgraph") - return await subgraph.ainvoke(state) + async def node(state: ParentState): + called.append("parent") + if state["counter"] < 1: + return Command( + goto="call_subgraph", update={"counter": state["counter"] + 1} + ) - async def node(state: ParentState): - called.append("parent") - if state["counter"] < 1: - return Command( - goto="call_subgraph", update={"counter": state["counter"] + 1} - ) + return {"foo": state["foo"] + "|" + "parent"} - return {"foo": state["foo"] + "|" + "parent"} + parent = ( + StateGraph(ParentState) + .add_node(call_subgraph) + .add_node(node) + .add_edge(START, "call_subgraph") + .add_edge("call_subgraph", "node") + .compile(checkpointer=async_checkpointer) + ) - parent = ( - StateGraph(ParentState) - .add_node(call_subgraph) - .add_node(node) - .add_edge(START, "call_subgraph") - .add_edge("call_subgraph", "node") - .compile(checkpointer=checkpointer) - ) + config = {"configurable": {"thread_id": "1"}} + assert await parent.ainvoke({"foo": "", "counter": 0}, config) == { + "foo": "", + "counter": 0, + "__interrupt__": [ + Interrupt( + value="Provide value", + resumable=True, + ns=[AnyStr("call_subgraph"), AnyStr("subnode_2:")], + ) + ], + } + assert await parent.ainvoke(Command(resume="bar"), config) == { + "foo": "subgraph_2", + "counter": 1, + "__interrupt__": [ + Interrupt( + value="Provide value", + resumable=True, + ns=[AnyStr("call_subgraph"), AnyStr("subnode_2")], + ) + ], + } + assert await parent.ainvoke(Command(resume="qux"), config) == { + "foo": "subgraph_2|parent", + "counter": 1, + } + assert called == [ + "call_subgraph", + "subnode_1", + "subnode_2", + "call_subgraph", + "subnode_2", + "parent", + "call_subgraph", + "subnode_1", + "subnode_2", + "call_subgraph", + "subnode_2", + "parent", + ] - config = {"configurable": {"thread_id": "1"}} - assert await parent.ainvoke({"foo": "", "counter": 0}, config) == { - "foo": "", - "counter": 0, - "__interrupt__": [ - Interrupt( - value="Provide value", - resumable=True, - ns=[AnyStr("call_subgraph"), AnyStr("subnode_2:")], - ) - ], - } - assert await parent.ainvoke(Command(resume="bar"), config) == { - "foo": "subgraph_2", - "counter": 1, - "__interrupt__": [ - Interrupt( - value="Provide value", - resumable=True, - ns=[AnyStr("call_subgraph"), AnyStr("subnode_2")], - ) - ], - } - assert await parent.ainvoke(Command(resume="qux"), config) == { - "foo": "subgraph_2|parent", - "counter": 1, - } - assert called == [ - "call_subgraph", - "subnode_1", - "subnode_2", - "call_subgraph", - "subnode_2", - "parent", - "call_subgraph", - "subnode_1", - "subnode_2", - "call_subgraph", - "subnode_2", - "parent", - ] - - # invoke parent again (new turn) - assert await parent.ainvoke({"foo": "meow", "counter": 0}, config) == { - "foo": "meow", - "counter": 0, - "__interrupt__": [ - Interrupt( - value="Provide value", - resumable=True, - ns=[AnyStr("call_subgraph"), AnyStr("subnode_2:")], - ) - ], - } - # confirm that we preserve the state values from the previous invocation - assert bar_values == [None, "barbaz", "quxbaz"] + # invoke parent again (new turn) + assert await parent.ainvoke({"foo": "meow", "counter": 0}, config) == { + "foo": "meow", + "counter": 0, + "__interrupt__": [ + Interrupt( + value="Provide value", + resumable=True, + ns=[AnyStr("call_subgraph"), AnyStr("subnode_2:")], + ) + ], + } + # confirm that we preserve the state values from the previous invocation + assert bar_values == [None, "barbaz", "quxbaz"] @NEEDS_CONTEXTVARS diff --git a/libs/prebuilt/tests/conftest.py b/libs/prebuilt/tests/conftest.py index 8ba0771c2..343856bd5 100644 --- a/libs/prebuilt/tests/conftest.py +++ b/libs/prebuilt/tests/conftest.py @@ -2,8 +2,6 @@ from collections.abc import AsyncIterator, Iterator from uuid import UUID import pytest -from langchain_core import __version__ as core_version -from packaging import version from pytest_mock import MockerFixture from langgraph.checkpoint.base import BaseCheckpointSaver @@ -31,11 +29,6 @@ from tests.conftest_store import ( pytest.register_assert_rewrite("tests.memory_assert") -# TODO: fix this once core is released -IS_LANGCHAIN_CORE_030_OR_GREATER = version.parse(core_version) >= version.parse( - "0.3.0.dev0" -) - @pytest.fixture def anyio_backend(): diff --git a/libs/prebuilt/tests/test_react_agent.py b/libs/prebuilt/tests/test_react_agent.py index 649b9dd6e..02581db22 100644 --- a/libs/prebuilt/tests/test_react_agent.py +++ b/libs/prebuilt/tests/test_react_agent.py @@ -55,7 +55,6 @@ from langgraph.store.memory import InMemoryStore from langgraph.types import Command, Interrupt, interrupt from langgraph.utils.config import get_stream_writer from tests.any_str import AnyStr -from tests.conftest import IS_LANGCHAIN_CORE_030_OR_GREATER from tests.messages import _AnyIdHumanMessage, _AnyIdToolMessage from tests.model import FakeToolCallingModel @@ -456,10 +455,6 @@ def test__infer_handled_types() -> None: _infer_handled_types(handler) -@pytest.mark.skipif( - not IS_LANGCHAIN_CORE_030_OR_GREATER, - reason="Pydantic v1 is required for this test to pass in langchain-core < 0.3", -) @pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) def test_react_agent_with_structured_response(version: str) -> None: class WeatherResponse(BaseModel): @@ -496,10 +491,6 @@ class CustomStatePydantic(AgentStatePydantic): user_name: Optional[str] = None -@pytest.mark.skipif( - not IS_LANGCHAIN_CORE_030_OR_GREATER, - reason="Langchain core 0.3.0 or greater is required", -) @pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) @pytest.mark.parametrize("state_schema", [CustomState, CustomStatePydantic]) def test_react_agent_update_state( @@ -565,10 +556,6 @@ def test_react_agent_update_state( assert tool_message.name == "get_user_name" -@pytest.mark.skipif( - not IS_LANGCHAIN_CORE_030_OR_GREATER, - reason="Langchain core 0.3.0 or greater is required", -) @pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) def test_react_agent_parallel_tool_calls( sync_checkpointer: BaseCheckpointSaver, version: str @@ -832,10 +819,6 @@ def test_create_react_agent_inject_vars( assert result["foo"] == 2 -@pytest.mark.skipif( - not IS_LANGCHAIN_CORE_030_OR_GREATER, - reason="Langchain core 0.3.0 or greater is required", -) def test_tool_node_inject_store() -> None: store = InMemoryStore() namespace = ("test",) diff --git a/libs/prebuilt/tests/test_tool_node.py b/libs/prebuilt/tests/test_tool_node.py index 0f2a908dc..b94094a23 100644 --- a/libs/prebuilt/tests/test_tool_node.py +++ b/libs/prebuilt/tests/test_tool_node.py @@ -18,7 +18,6 @@ from langgraph.errors import NodeInterrupt from langgraph.prebuilt import ToolNode from langgraph.prebuilt.tool_node import TOOL_CALL_ERROR_TEMPLATE from langgraph.types import Command, Send -from tests.conftest import IS_LANGCHAIN_CORE_030_OR_GREATER pytestmark = pytest.mark.anyio @@ -492,10 +491,6 @@ def test_tool_node_node_interrupt(): assert exc_info.value == "foo" -@pytest.mark.skipif( - not IS_LANGCHAIN_CORE_030_OR_GREATER, - reason="Langchain core 0.3.0 or greater is required", -) @pytest.mark.parametrize("input_type", ["dict", "tool_calls"]) async def test_tool_node_command(input_type: str): from langchain_core.tools.base import InjectedToolCallId @@ -797,10 +792,6 @@ async def test_tool_node_command(input_type: str): ) == [Command(update={"messages": []}, graph=Command.PARENT)] -@pytest.mark.skipif( - not IS_LANGCHAIN_CORE_030_OR_GREATER, - reason="Langchain core 0.3.0 or greater is required", -) async def test_tool_node_command_list_input(): from langchain_core.tools.base import InjectedToolCallId