Note
@@ -27,101 +42,162 @@ The LangGraph CLI requires a JSON configuration file with the following keys:Prerequisites
\n", + "\n", + " This guide assumes familiarity with the following:\n", + "
Set up LangSmith for LangGraph development
\n", + "\n", + " Sign up for LangSmith to quickly spot issues and improve the performance of your LangGraph projects. LangSmith lets you use trace data to debug, test, and monitor your LLM apps built with LangGraph — read more about how to get started here. \n", + "
\n", + "Prerequisites
\n", @@ -17,25 +17,27 @@ " This guide assumes familiarity with the following:\n", "Prerequisites
\n", + "\n", + " This guide assumes familiarity with the following:\n", + "
Note
\n", "\n",
- " The first thing you do when you define a graph is define the State of the graph. The State consists of the schema of the graph as well as reducer functions which specify how to apply updates to the state. In our example State is a TypedDict with a single key: messages. The messages key is annotated with the add_messages reducer function, which tells LangGraph to append new messages to the existing list, rather than overwriting it. State keys without an annotation will be overwritten by each update, storing the most recent value. Check out this conceptual guide to learn more about state, reducers and other low-level concepts.\n",
+ " The first thing you do when you define a graph is define the State of the graph. The State consists of the schema of the graph as well as reducer functions which specify how to apply updates to the state. In our example State is a TypedDict with a single key: messages. The messages key is annotated with the add_messages reducer function, which tells LangGraph to append new messages to the existing list, rather than overwriting it. State keys without an annotation will be overwritten by each update, storing the most recent value. Check out this conceptual guide to learn more about state, reducers and other low-level concepts.\n",
"
__start__
]):::first + prepare(prepare) + tool_two_slow(tool_two_slow) + tool_two_fast(tool_two_fast) + finish(finish) + __end__([__end__
]):::last + __start__ --> prepare; + finish --> __end__; + prepare -.-> tool_two_slow; + tool_two_slow --> finish; + prepare -.-> tool_two_fast; + tool_two_fast --> finish; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- # name: test_branch_then[memory] ''' graph TD; @@ -695,6 +729,281 @@ ''' # --- +# name: test_conditional_graph[duckdb] + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableAssign" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "runnable", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "parents": {}, + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[duckdb].1 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + + ''' +# --- +# name: test_conditional_graph[duckdb].2 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + tools(tools__end__
]):::last + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_conditional_graph[duckdb].3 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": "__start__" + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableAssign" + ], + "name": "agent" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langgraph", + "utils", + "runnable", + "RunnableCallable" + ], + "name": "tools" + }, + "metadata": { + "parents": {}, + "version": 2, + "variant": "b" + } + }, + { + "id": "__end__", + "type": "schema", + "data": "__end__" + } + ], + "edges": [ + { + "source": "__start__", + "target": "agent" + }, + { + "source": "tools", + "target": "agent" + }, + { + "source": "agent", + "target": "tools", + "data": "continue", + "conditional": true + }, + { + "source": "agent", + "target": "__end__", + "data": "exit", + "conditional": true + } + ] + } + ''' +# --- +# name: test_conditional_graph[duckdb].4 + ''' + graph TD; + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + + ''' +# --- +# name: test_conditional_graph[duckdb].5 + dict({ + 'edges': list([ + dict({ + 'source': '__start__', + 'target': 'agent', + }), + dict({ + 'source': 'tools', + 'target': 'agent', + }), + dict({ + 'conditional': True, + 'data': 'continue', + 'source': 'agent', + 'target': 'tools', + }), + dict({ + 'conditional': True, + 'data': 'exit', + 'source': 'agent', + 'target': '__end__', + }), + ]), + 'nodes': list([ + dict({ + 'data': '__start__', + 'id': '__start__', + 'type': 'schema', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langchain', + 'schema', + 'runnable', + 'RunnableAssign', + ]), + 'name': 'agent', + }), + 'id': 'agent', + 'metadata': dict({ + '__interrupt': 'after', + }), + 'type': 'runnable', + }), + dict({ + 'data': dict({ + 'id': list([ + 'langgraph', + 'utils', + 'runnable', + 'RunnableCallable', + ]), + 'name': 'tools', + }), + 'id': 'tools', + 'metadata': dict({ + 'parents': dict({ + }), + 'variant': 'b', + 'version': 2, + }), + 'type': 'runnable', + }), + dict({ + 'data': '__end__', + 'id': '__end__', + 'type': 'schema', + }), + ]), + }) +# --- +# name: test_conditional_graph[duckdb].6 + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent__end__
]):::last + __start__ --> agent; + tools --> agent; + agent -. continue .-> tools; + agent -. exit .-> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- # name: test_conditional_graph[memory] ''' { @@ -806,65 +1115,16 @@ "data": "__start__" }, { - "id": 1, - "type": "schema", - "data": "Parallel__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[postgres] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[postgres_pipe] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[postgres_pool] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[sqlite] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- # name: test_simple_multi_edge ''' graph TD; @@ -4929,6 +5195,24 @@ ''' # --- +# name: test_start_branch_then[duckdb] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + tool_two_slow(tool_two_slow) + tool_two_fast(tool_two_fast) + __end__([__end__
]):::last + __start__ -.-> tool_two_slow; + tool_two_slow --> __end__; + __start__ -.-> tool_two_fast; + tool_two_fast --> __end__; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- # name: test_start_branch_then[memory] ''' %%{init: {'flowchart': {'curve': 'linear'}}}%% @@ -5046,6 +5330,31 @@ # name: test_state_graph_w_config_inherited_state_keys.2 '{"$defs": {"AgentAction": {"description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"anyOf": [{"type": "string"}, {"type": "object"}], "title": "Tool Input"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentAction", "default": "AgentAction", "enum": ["AgentAction"], "title": "Type", "type": "string"}}, "required": ["tool", "tool_input", "log"], "title": "AgentAction", "type": "object"}, "AgentFinish": {"description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"const": "AgentFinish", "default": "AgentFinish", "enum": ["AgentFinish"], "title": "Type", "type": "string"}}, "required": ["return_values", "log"], "title": "AgentFinish", "type": "object"}}, "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"anyOf": [{"$ref": "#/$defs/AgentAction"}, {"$ref": "#/$defs/AgentFinish"}, {"type": "null"}], "default": null, "title": "Agent Outcome"}, "intermediate_steps": {"default": null, "items": {"maxItems": 2, "minItems": 2, "prefixItems": [{"$ref": "#/$defs/AgentAction"}, {"type": "string"}], "type": "array"}, "title": "Intermediate Steps", "type": "array"}}, "required": ["input"], "title": "LangGraphOutput", "type": "object"}' # --- +# name: test_weather_subgraph[duckdb] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + router_node(router_node) + normal_llm_node(normal_llm_node) + weather_graph_model_node(model_node) + weather_graph_weather_node(weather_node__end__
]):::last + __start__ --> router_node; + normal_llm_node --> __end__; + weather_graph_weather_node --> __end__; + router_node -.-> normal_llm_node; + router_node -.-> weather_graph_model_node; + router_node -.-> __end__; + subgraph weather_graph + weather_graph_model_node --> weather_graph_weather_node; + end + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- # name: test_weather_subgraph[memory] ''' %%{init: {'flowchart': {'curve': 'linear'}}}%% diff --git a/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr b/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr index 0acb4d0a8..46916c7a4 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr @@ -221,19 +221,6 @@ +---------+ ''' # --- -# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2 - ''' - graph TD; - __start__ --> rewrite_query; - analyzer_one --> retriever_one; - qa --> __end__; - retriever_one --> qa; - retriever_two --> qa; - rewrite_query --> analyzer_one; - rewrite_query -.-> retriever_two; - - ''' -# --- # name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2.1 dict({ '$defs': dict({ @@ -342,6 +329,127 @@ 'type': 'object', }) # --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[duckdb_aio] + ''' + graph TD; + __start__ --> rewrite_query; + analyzer_one --> retriever_one; + qa --> __end__; + retriever_one --> qa; + retriever_two --> qa; + rewrite_query --> analyzer_one; + rewrite_query -.-> retriever_two; + + ''' +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[duckdb_aio].1 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'answer': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'default': None, + 'title': 'Answer', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + 'docs', + ]), + 'title': 'State', + 'type': 'object', + }) +# --- +# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[duckdb_aio].2 + dict({ + '$defs': dict({ + 'InnerObject': dict({ + 'properties': dict({ + 'yo': dict({ + 'title': 'Yo', + 'type': 'integer', + }), + }), + 'required': list([ + 'yo', + ]), + 'title': 'InnerObject', + 'type': 'object', + }), + }), + 'properties': dict({ + 'answer': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'default': None, + 'title': 'Answer', + }), + 'docs': dict({ + 'items': dict({ + 'type': 'string', + }), + 'title': 'Docs', + 'type': 'array', + }), + 'inner': dict({ + '$ref': '#/$defs/InnerObject', + }), + 'query': dict({ + 'title': 'Query', + 'type': 'string', + }), + }), + 'required': list([ + 'query', + 'inner', + 'docs', + ]), + 'title': 'State', + 'type': 'object', + }) +# --- # name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[memory] ''' graph TD; @@ -1194,6 +1302,106 @@ +---------+ ''' # --- +# name: test_send_react_interrupt_control[memory] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[postgres_aio] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[postgres_aio_pipe] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[postgres_aio_pool] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_send_react_interrupt_control[sqlite_aio] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + agent(agent) + foo([foo]):::last + __start__ --> agent; + agent -.-> foo; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- +# name: test_weather_subgraph[duckdb_aio] + ''' + %%{init: {'flowchart': {'curve': 'linear'}}}%% + graph TD; + __start__([__start__
]):::first + router_node(router_node) + normal_llm_node(normal_llm_node) + weather_graph_model_node(model_node) + weather_graph_weather_node(weather_node__end__
]):::last + __start__ --> router_node; + normal_llm_node --> __end__; + weather_graph_weather_node --> __end__; + router_node -.-> normal_llm_node; + router_node -.-> weather_graph_model_node; + router_node -.-> __end__; + subgraph weather_graph + weather_graph_model_node --> weather_graph_weather_node; + end + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc + + ''' +# --- # name: test_weather_subgraph[memory] ''' %%{init: {'flowchart': {'curve': 'linear'}}}%% diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index a58923f14..eae7694ff 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -11,14 +11,18 @@ from psycopg_pool import AsyncConnectionPool, ConnectionPool from pytest_mock import MockerFixture from langgraph.checkpoint.base import BaseCheckpointSaver +from langgraph.checkpoint.duckdb import DuckDBSaver +from langgraph.checkpoint.duckdb.aio import AsyncDuckDBSaver from langgraph.checkpoint.postgres import PostgresSaver from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver from langgraph.checkpoint.sqlite import SqliteSaver from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver from langgraph.store.base import BaseStore +from langgraph.store.duckdb import AsyncDuckDBStore, DuckDBStore from langgraph.store.memory import InMemoryStore from langgraph.store.postgres import AsyncPostgresStore, PostgresStore -from tests.memory_assert import MemorySaverAssertImmutable + +pytest.register_assert_rewrite("tests.memory_assert") DEFAULT_POSTGRES_URI = "postgres://postgres:postgres@localhost:5442/" # TODO: fix this once core is released @@ -46,6 +50,8 @@ def deterministic_uuids(mocker: MockerFixture) -> MockerFixture: @pytest.fixture(scope="function") def checkpointer_memory(): + from tests.memory_assert import MemorySaverAssertImmutable + yield MemorySaverAssertImmutable() @@ -61,6 +67,20 @@ async def _checkpointer_sqlite_aio(): yield checkpointer +@pytest.fixture(scope="function") +def checkpointer_duckdb(): + with DuckDBSaver.from_conn_string(":memory:") as checkpointer: + checkpointer.setup() + yield checkpointer + + +@asynccontextmanager +async def _checkpointer_duckdb_aio(): + async with AsyncDuckDBSaver.from_conn_string(":memory:") as checkpointer: + await checkpointer.setup() + yield checkpointer + + @pytest.fixture(scope="function") def checkpointer_postgres(): database = f"test_{uuid4().hex[:16]}" @@ -208,10 +228,15 @@ async def awith_checkpointer( if checkpointer_name is None: yield None elif checkpointer_name == "memory": + from tests.memory_assert import MemorySaverAssertImmutable + yield MemorySaverAssertImmutable() elif checkpointer_name == "sqlite_aio": async with _checkpointer_sqlite_aio() as checkpointer: yield checkpointer + elif checkpointer_name == "duckdb_aio": + async with _checkpointer_duckdb_aio() as checkpointer: + yield checkpointer elif checkpointer_name == "postgres_aio": async with _checkpointer_postgres_aio() as checkpointer: yield checkpointer @@ -247,6 +272,13 @@ async def _store_postgres_aio(): await conn.execute(f"DROP DATABASE {database}") +@asynccontextmanager +async def _store_duckdb_aio(): + async with AsyncDuckDBStore.from_conn_string(":memory:") as store: + await store.setup() + yield store + + @pytest.fixture(scope="function") def store_postgres(): database = f"test_{uuid4().hex[:16]}" @@ -264,6 +296,13 @@ def store_postgres(): conn.execute(f"DROP DATABASE {database}") +@pytest.fixture(scope="function") +def store_duckdb(): + with DuckDBStore.from_conn_string(":memory:") as store: + store.setup() + yield store + + @pytest.fixture(scope="function") def store_in_memory(): yield InMemoryStore() @@ -278,6 +317,9 @@ async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]: elif store_name == "postgres_aio": async with _store_postgres_aio() as store: yield store + elif store_name == "duckdb_aio": + async with _store_duckdb_aio() as store: + yield store else: raise NotImplementedError(f"Unknown store {store_name}") @@ -300,5 +342,5 @@ ALL_CHECKPOINTERS_ASYNC_PLUS_NONE = [ *ALL_CHECKPOINTERS_ASYNC, None, ] -ALL_STORES_SYNC = ["in_memory", "postgres"] -ALL_STORES_ASYNC = ["in_memory", "postgres_aio"] +ALL_STORES_SYNC = ["in_memory", "postgres", "duckdb"] +ALL_STORES_ASYNC = ["in_memory", "postgres_aio", "duckdb_aio"] diff --git a/libs/langgraph/tests/memory_assert.py b/libs/langgraph/tests/memory_assert.py index 6b44051f7..0a9f13a47 100644 --- a/libs/langgraph/tests/memory_assert.py +++ b/libs/langgraph/tests/memory_assert.py @@ -1,5 +1,8 @@ import asyncio +import os +import tempfile from collections import defaultdict +from functools import partial from typing import Any, Optional from langchain_core.runnables import RunnableConfig @@ -12,7 +15,7 @@ from langgraph.checkpoint.base import ( SerializerProtocol, copy_checkpoint, ) -from langgraph.checkpoint.memory import MemorySaver +from langgraph.checkpoint.memory import MemorySaver, PersistentDict class NoopSerializer(SerializerProtocol): @@ -32,9 +35,13 @@ class MemorySaverAssertImmutable(MemorySaver): serde: Optional[SerializerProtocol] = None, put_sleep: Optional[float] = None, ) -> None: - super().__init__(serde=serde) + _, filename = tempfile.mkstemp() + super().__init__( + serde=serde, factory=partial(PersistentDict, filename=filename) + ) self.storage_for_copies = defaultdict(lambda: defaultdict(dict)) self.put_sleep = put_sleep + self.stack.callback(os.remove, filename) def put( self, diff --git a/libs/langgraph/tests/test_algo.py b/libs/langgraph/tests/test_algo.py index 4e259f29e..9d6ec5942 100644 --- a/libs/langgraph/tests/test_algo.py +++ b/libs/langgraph/tests/test_algo.py @@ -11,13 +11,21 @@ def test_prepare_next_tasks() -> None: with ChannelsManager({}, checkpoint, config) as (channels, managed): assert ( prepare_next_tasks( - checkpoint, processes, channels, managed, config, 0, for_execution=False + checkpoint, + {}, + processes, + channels, + managed, + config, + 0, + for_execution=False, ) == {} ) assert ( prepare_next_tasks( checkpoint, + {}, processes, channels, managed, diff --git a/libs/langgraph/tests/test_prebuilt.py b/libs/langgraph/tests/test_prebuilt.py index c1176f365..a6655a451 100644 --- a/libs/langgraph/tests/test_prebuilt.py +++ b/libs/langgraph/tests/test_prebuilt.py @@ -1,5 +1,6 @@ import dataclasses import json +from functools import partial from typing import ( Annotated, Any, @@ -28,25 +29,41 @@ from langchain_core.messages import ( ) from langchain_core.outputs import ChatGeneration, ChatResult from langchain_core.runnables import Runnable, RunnableLambda -from langchain_core.tools import BaseTool +from langchain_core.tools import BaseTool, ToolException from langchain_core.tools import tool as dec_tool -from pydantic import BaseModel +from pydantic import BaseModel, ValidationError from pydantic.v1 import BaseModel as BaseModelV1 +from pydantic.v1 import ValidationError as ValidationErrorV1 from typing_extensions import TypedDict from langgraph.checkpoint.base import BaseCheckpointSaver -from langgraph.graph import START, MessagesState, StateGraph -from langgraph.prebuilt import ToolNode, ValidationNode, create_react_agent -from langgraph.prebuilt.tool_node import InjectedState, InjectedStore +from langgraph.checkpoint.memory import MemorySaver +from langgraph.errors import NodeInterrupt +from langgraph.graph import START, MessagesState, StateGraph, add_messages +from langgraph.prebuilt import ( + ToolNode, + ValidationNode, + create_react_agent, + tools_condition, +) +from langgraph.prebuilt.chat_agent_executor import _validate_chat_history +from langgraph.prebuilt.tool_node import ( + TOOL_CALL_ERROR_TEMPLATE, + InjectedState, + InjectedStore, + _get_state_args, + _infer_handled_types, +) from langgraph.store.base import BaseStore from langgraph.store.memory import InMemoryStore +from langgraph.types import Interrupt from tests.conftest import ( ALL_CHECKPOINTERS_ASYNC, ALL_CHECKPOINTERS_SYNC, IS_LANGCHAIN_CORE_030_OR_GREATER, awith_checkpointer, ) -from tests.messages import _AnyIdHumanMessage +from tests.messages import _AnyIdHumanMessage, _AnyIdToolMessage pytestmark = pytest.mark.anyio @@ -141,6 +158,7 @@ def test_no_modifier(request: pytest.FixtureRequest, checkpointer_name: str) -> "source": "loop", "writes": {"agent": {"messages": [AIMessage(content="hi?", id="0")]}}, "step": 1, + "thread_id": "123", } assert saved.pending_writes == [] @@ -172,6 +190,7 @@ async def test_no_modifier_async(checkpointer_name: str) -> None: "source": "loop", "writes": {"agent": {"messages": [AIMessage(content="hi?", id="0")]}}, "step": 1, + "thread_id": "123", } assert saved.pending_writes == [] @@ -362,32 +381,172 @@ def test_model_with_tools(tool_style: str): create_react_agent(model.bind_tools([tool1]), [tool2]) +def test__validate_messages(): + # empty input + _validate_chat_history([]) + + # single human message + _validate_chat_history( + [ + HumanMessage(content="What's the weather?"), + ] + ) + + # human + AI + _validate_chat_history( + [ + HumanMessage(content="What's the weather?"), + AIMessage(content="The weather is sunny and 75°F."), + ] + ) + + # Answered tool calls + _validate_chat_history( + [ + HumanMessage(content="What's the weather?"), + AIMessage( + content="Let me check that for you.", + tool_calls=[{"id": "call1", "name": "get_weather", "args": {}}], + ), + ToolMessage(content="Sunny, 75°F", tool_call_id="call1"), + AIMessage(content="The weather is sunny and 75°F."), + ] + ) + + # Unanswered tool calls + with pytest.raises(ValueError): + _validate_chat_history( + [ + AIMessage( + content="I'll check that for you.", + tool_calls=[ + {"id": "call1", "name": "get_weather", "args": {}}, + {"id": "call2", "name": "get_time", "args": {}}, + ], + ) + ] + ) + + with pytest.raises(ValueError): + _validate_chat_history( + [ + HumanMessage(content="What's the weather and time?"), + AIMessage( + content="I'll check that for you.", + tool_calls=[ + {"id": "call1", "name": "get_weather", "args": {}}, + {"id": "call2", "name": "get_time", "args": {}}, + ], + ), + ToolMessage(content="Sunny, 75°F", tool_call_id="call1"), + AIMessage( + content="The weather is sunny and 75°F. Let me check the time." + ), + ] + ) + + +def test__infer_handled_types() -> None: + def handle(e): # type: ignore + return "" + + def handle2(e: Exception) -> str: + return "" + + def handle3(e: Union[ValueError, ToolException]) -> str: + return "" + + class Handler: + def handle(self, e: ValueError) -> str: + return "" + + handle4 = Handler().handle + + def handle5(e: Union[Union[TypeError, ValueError], ToolException]): + return "" + + expected: tuple = (Exception,) + actual = _infer_handled_types(handle) + assert expected == actual + + expected = (Exception,) + actual = _infer_handled_types(handle2) + assert expected == actual + + expected = (ValueError, ToolException) + actual = _infer_handled_types(handle3) + assert expected == actual + + expected = (ValueError,) + actual = _infer_handled_types(handle4) + assert expected == actual + + expected = (TypeError, ValueError, ToolException) + actual = _infer_handled_types(handle5) + assert expected == actual + + with pytest.raises(ValueError): + + def handler(e: str): + return "" + + _infer_handled_types(handler) + + with pytest.raises(ValueError): + + def handler(e: list[Exception]): + return "" + + _infer_handled_types(handler) + + with pytest.raises(ValueError): + + def handler(e: Union[str, int]): + return "" + + _infer_handled_types(handler) + + +# tools for testing Too +def tool1(some_val: int, some_other_val: str) -> str: + """Tool 1 docstring.""" + if some_val == 0: + raise ValueError("Test error") + return f"{some_val} - {some_other_val}" + + +async def tool2(some_val: int, some_other_val: str) -> str: + """Tool 2 docstring.""" + if some_val == 0: + raise ToolException("Test error") + return f"tool2: {some_val} - {some_other_val}" + + +async def tool3(some_val: int, some_other_val: str) -> str: + """Tool 3 docstring.""" + return [ + {"key_1": some_val, "key_2": "foo"}, + {"key_1": some_other_val, "key_2": "baz"}, + ] + + +async def tool4(some_val: int, some_other_val: str) -> str: + """Tool 4 docstring.""" + return [ + {"type": "image_url", "image_url": {"url": "abdc"}}, + ] + + +@dec_tool +def tool5(some_val: int): + """Tool 5 docstring.""" + raise ToolException("Test error") + + +tool5.handle_tool_error = "foo" + + async def test_tool_node(): - def tool1(some_val: int, some_other_val: str) -> str: - """Tool 1 docstring.""" - if some_val == 0: - raise ValueError("Test error") - return f"{some_val} - {some_other_val}" - - async def tool2(some_val: int, some_other_val: str) -> str: - """Tool 2 docstring.""" - if some_val == 0: - raise ValueError("Test error") - return f"tool2: {some_val} - {some_other_val}" - - async def tool3(some_val: int, some_other_val: str) -> str: - """Tool 3 docstring.""" - return [ - {"key_1": some_val, "key_2": "foo"}, - {"key_1": some_other_val, "key_2": "baz"}, - ] - - async def tool4(some_val: int, some_other_val: str) -> str: - """Tool 4 docstring.""" - return [ - {"type": "image_url", "image_url": {"url": "abdc"}}, - ] - result = ToolNode([tool1]).invoke( { "messages": [ @@ -410,31 +569,6 @@ async def test_tool_node(): assert tool_message.content == "1 - foo" assert tool_message.tool_call_id == "some 0" - result_error = ToolNode([tool1]).invoke( - { - "messages": [ - AIMessage( - "hi?", - tool_calls=[ - { - "name": "tool1", - "args": {"some_val": 0, "some_other_val": "foo"}, - "id": "some 0", - } - ], - ) - ] - } - ) - - tool_message: ToolMessage = result_error["messages"][-1] - assert tool_message.type == "tool" - assert ( - tool_message.content - == f"Error: {repr(ValueError('Test error'))}\n Please fix your mistakes." - ) - assert tool_message.tool_call_id == "some 0" - result2 = await ToolNode([tool2]).ainvoke( { "messages": [ @@ -451,11 +585,232 @@ async def test_tool_node(): ] } ) + tool_message: ToolMessage = result2["messages"][-1] assert tool_message.type == "tool" assert tool_message.content == "tool2: 2 - bar" - with pytest.raises(ValueError): + # list of dicts tool content + result3 = await ToolNode([tool3]).ainvoke( + { + "messages": [ + AIMessage( + "hi?", + tool_calls=[ + { + "name": "tool3", + "args": {"some_val": 2, "some_other_val": "bar"}, + "id": "some 2", + } + ], + ) + ] + } + ) + tool_message: ToolMessage = result3["messages"][-1] + assert tool_message.type == "tool" + assert ( + tool_message.content + == '[{"key_1": 2, "key_2": "foo"}, {"key_1": "bar", "key_2": "baz"}]' + ) + assert tool_message.tool_call_id == "some 2" + + # list of content blocks tool content + result4 = await ToolNode([tool4]).ainvoke( + { + "messages": [ + AIMessage( + "hi?", + tool_calls=[ + { + "name": "tool4", + "args": {"some_val": 2, "some_other_val": "bar"}, + "id": "some 3", + } + ], + ) + ] + } + ) + tool_message: ToolMessage = result4["messages"][-1] + assert tool_message.type == "tool" + assert tool_message.content == [{"type": "image_url", "image_url": {"url": "abdc"}}] + assert tool_message.tool_call_id == "some 3" + + +async def test_tool_node_error_handling(): + def handle_all(e: Union[ValueError, ToolException, ValidationError]): + return TOOL_CALL_ERROR_TEMPLATE.format(error=repr(e)) + + # test catching all exceptions, via: + # - handle_tool_errors = True + # - passing a tuple of all exceptions + # - passing a callable with all exceptions in the signature + for handle_tool_errors in ( + True, + (ValueError, ToolException, ValidationError), + handle_all, + ): + result_error = await ToolNode( + [tool1, tool2, tool3], handle_tool_errors=handle_tool_errors + ).ainvoke( + { + "messages": [ + AIMessage( + "hi?", + tool_calls=[ + { + "name": "tool1", + "args": {"some_val": 0, "some_other_val": "foo"}, + "id": "some id", + }, + { + "name": "tool2", + "args": {"some_val": 0, "some_other_val": "bar"}, + "id": "some other id", + }, + { + "name": "tool3", + "args": {"some_val": 0}, + "id": "another id", + }, + ], + ) + ] + } + ) + + assert all(m.type == "tool" for m in result_error["messages"]) + assert all(m.status == "error" for m in result_error["messages"]) + assert ( + result_error["messages"][0].content + == f"Error: {repr(ValueError('Test error'))}\n Please fix your mistakes." + ) + assert ( + result_error["messages"][1].content + == f"Error: {repr(ToolException('Test error'))}\n Please fix your mistakes." + ) + assert ( + "ValidationError" in result_error["messages"][2].content + or "validation error" in result_error["messages"][2].content + ) + + assert result_error["messages"][0].tool_call_id == "some id" + assert result_error["messages"][1].tool_call_id == "some other id" + assert result_error["messages"][2].tool_call_id == "another id" + + +async def test_tool_node_error_handling_callable(): + def handle_value_error(e: ValueError): + return "Value error" + + def handle_tool_exception(e: ToolException): + return "Tool exception" + + for handle_tool_errors in ("Value error", handle_value_error): + result_error = await ToolNode( + [tool1], handle_tool_errors=handle_tool_errors + ).ainvoke( + { + "messages": [ + AIMessage( + "hi?", + tool_calls=[ + { + "name": "tool1", + "args": {"some_val": 0, "some_other_val": "foo"}, + "id": "some id", + }, + ], + ) + ] + } + ) + tool_message: ToolMessage = result_error["messages"][-1] + assert tool_message.type == "tool" + assert tool_message.status == "error" + assert tool_message.content == "Value error" + + # test raising for an unhandled exception, via: + # - passing a tuple of all exceptions + # - passing a callable with all exceptions in the signature + for handle_tool_errors in ((ValueError,), handle_value_error): + with pytest.raises(ToolException) as exc_info: + await ToolNode( + [tool1, tool2], handle_tool_errors=handle_tool_errors + ).ainvoke( + { + "messages": [ + AIMessage( + "hi?", + tool_calls=[ + { + "name": "tool1", + "args": {"some_val": 0, "some_other_val": "foo"}, + "id": "some id", + }, + { + "name": "tool2", + "args": {"some_val": 0, "some_other_val": "bar"}, + "id": "some other id", + }, + ], + ) + ] + } + ) + assert str(exc_info.value) == "Test error" + + for handle_tool_errors in ((ToolException,), handle_tool_exception): + with pytest.raises(ValueError) as exc_info: + await ToolNode( + [tool1, tool2], handle_tool_errors=handle_tool_errors + ).ainvoke( + { + "messages": [ + AIMessage( + "hi?", + tool_calls=[ + { + "name": "tool1", + "args": {"some_val": 0, "some_other_val": "foo"}, + "id": "some id", + }, + { + "name": "tool2", + "args": {"some_val": 0, "some_other_val": "bar"}, + "id": "some other id", + }, + ], + ) + ] + } + ) + assert str(exc_info.value) == "Test error" + + +async def test_tool_node_handle_tool_errors_false(): + with pytest.raises(ValueError) as exc_info: + ToolNode([tool1], handle_tool_errors=False).invoke( + { + "messages": [ + AIMessage( + "hi?", + tool_calls=[ + { + "name": "tool1", + "args": {"some_val": 0, "some_other_val": "foo"}, + "id": "some id", + } + ], + ) + ] + } + ) + + assert str(exc_info.value) == "Test error" + + with pytest.raises(ToolException): await ToolNode([tool2], handle_tool_errors=False).ainvoke( { "messages": [ @@ -465,7 +820,7 @@ async def test_tool_node(): { "name": "tool2", "args": {"some_val": 0, "some_other_val": "bar"}, - "id": "some 1", + "id": "some id", } ], ) @@ -473,7 +828,57 @@ async def test_tool_node(): } ) - # incorrect tool name + assert str(exc_info.value) == "Test error" + + # test validation errors get raised if handle_tool_errors is False + with pytest.raises((ValidationError, ValidationErrorV1)): + ToolNode([tool1], handle_tool_errors=False).invoke( + { + "messages": [ + AIMessage( + "hi?", + tool_calls=[ + { + "name": "tool1", + "args": {"some_val": 0}, + "id": "some id", + } + ], + ) + ] + } + ) + + +def test_tool_node_individual_tool_error_handling(): + # test error handling on individual tools (and that it overrides overall error handling!) + result_individual_tool_error_handler = ToolNode( + [tool5], handle_tool_errors="bar" + ).invoke( + { + "messages": [ + AIMessage( + "hi?", + tool_calls=[ + { + "name": "tool5", + "args": {"some_val": 0}, + "id": "some 0", + } + ], + ) + ] + } + ) + + tool_message: ToolMessage = result_individual_tool_error_handler["messages"][-1] + assert tool_message.type == "tool" + assert tool_message.status == "error" + assert tool_message.content == "foo" + assert tool_message.tool_call_id == "some 0" + + +def test_tool_node_incorrect_tool_name(): result_incorrect_name = ToolNode([tool1, tool2]).invoke( { "messages": [ @@ -490,60 +895,94 @@ async def test_tool_node(): ] } ) + tool_message: ToolMessage = result_incorrect_name["messages"][-1] assert tool_message.type == "tool" + assert tool_message.status == "error" assert ( tool_message.content == "Error: tool3 is not a valid tool, try one of [tool1, tool2]." ) assert tool_message.tool_call_id == "some 0" - # list of dicts tool content - result3 = await ToolNode([tool3]).ainvoke( - { - "messages": [ - AIMessage( - "hi?", - tool_calls=[ - { - "name": "tool3", - "args": {"some_val": 2, "some_other_val": "bar"}, - "id": "some 0", - } - ], - ) - ] - } - ) - tool_message: ToolMessage = result3["messages"][-1] - assert tool_message.type == "tool" - assert ( - tool_message.content - == '[{"key_1": 2, "key_2": "foo"}, {"key_1": "bar", "key_2": "baz"}]' - ) - assert tool_message.tool_call_id == "some 0" - # list of content blocks tool content - result4 = await ToolNode([tool4]).ainvoke( - { - "messages": [ - AIMessage( - "hi?", - tool_calls=[ - { - "name": "tool4", - "args": {"some_val": 2, "some_other_val": "bar"}, - "id": "some 0", - } - ], - ) - ] - } +def test_tool_node_node_interrupt(): + def tool_normal(some_val: int) -> str: + """Tool docstring.""" + return "normal" + + def tool_interrupt(some_val: int) -> str: + """Tool docstring.""" + raise NodeInterrupt("foo") + + def handle(e: NodeInterrupt): + return "handled" + + for handle_tool_errors in (True, (NodeInterrupt,), "handled", handle, False): + node = ToolNode([tool_interrupt], handle_tool_errors=handle_tool_errors) + with pytest.raises(NodeInterrupt) as exc_info: + node.invoke( + { + "messages": [ + AIMessage( + "hi?", + tool_calls=[ + { + "name": "tool_interrupt", + "args": {"some_val": 0}, + "id": "some 0", + } + ], + ) + ] + } + ) + assert exc_info.value == "foo" + + # test inside react agent + model = FakeToolCallingModel( + tool_calls=[ + [ + ToolCall(name="tool_interrupt", args={"some_val": 0}, id="1"), + ToolCall(name="tool_normal", args={"some_val": 1}, id="2"), + ], + [], + ] ) - tool_message: ToolMessage = result4["messages"][-1] - assert tool_message.type == "tool" - assert tool_message.content == [{"type": "image_url", "image_url": {"url": "abdc"}}] - assert tool_message.tool_call_id == "some 0" + checkpointer = MemorySaver() + config = {"configurable": {"thread_id": "1"}} + agent = create_react_agent( + model, [tool_interrupt, tool_normal], checkpointer=checkpointer + ) + result = agent.invoke({"messages": [HumanMessage("hi?")]}, config) + assert result["messages"] == [ + _AnyIdHumanMessage( + content="hi?", + ), + AIMessage( + content="hi?", + id="0", + tool_calls=[ + { + "name": "tool_interrupt", + "args": {"some_val": 0}, + "id": "1", + "type": "tool_call", + }, + { + "name": "tool_normal", + "args": {"some_val": 1}, + "id": "2", + "type": "tool_call", + }, + ], + ), + ] + state = agent.get_state(config) + assert state.next == ("tools",) + task = state.tasks[0] + assert task.name == "tools" + assert task.interrupts == (Interrupt(value="foo", when="during"),) def my_function(some_val: int, some_other_val: str) -> str: @@ -826,6 +1265,47 @@ def test_tool_node_ensure_utf8() -> None: assert outputs[0].content == json.dumps(data, ensure_ascii=False) +def test_tool_node_messages_key() -> None: + @dec_tool + def add(a: int, b: int): + """Adds a and b.""" + return a + b + + model = FakeToolCallingModel( + tool_calls=[[ToolCall(name=add.name, args={"a": 1, "b": 2}, id="test_id")]] + ) + + class State(TypedDict): + subgraph_messages: Annotated[list[AnyMessage], add_messages] + + def call_model(state: State): + response = model.invoke(state["subgraph_messages"]) + model.tool_calls = [] + return {"subgraph_messages": response} + + builder = StateGraph(State) + builder.add_node("agent", call_model) + builder.add_node("tools", ToolNode([add], messages_key="subgraph_messages")) + builder.add_conditional_edges( + "agent", partial(tools_condition, messages_key="subgraph_messages") + ) + builder.add_edge(START, "agent") + builder.add_edge("tools", "agent") + + graph = builder.compile() + result = graph.invoke({"subgraph_messages": [HumanMessage(content="hi")]}) + assert result["subgraph_messages"] == [ + _AnyIdHumanMessage(content="hi"), + AIMessage( + content="hi", + id="0", + tool_calls=[ToolCall(name=add.name, args={"a": 1, "b": 2}, id="test_id")], + ), + _AnyIdToolMessage(content="3", name=add.name, tool_call_id="test_id"), + AIMessage(content="hi-hi-3", id="1"), + ] + + async def test_return_direct() -> None: @dec_tool(return_direct=True) def tool_return_direct(input: str) -> str: @@ -921,3 +1401,18 @@ async def test_return_direct() -> None: id=result["messages"][3].id, ), ] + + +def test__get_state_args() -> None: + class Schema1(BaseModel): + a: Annotated[str, InjectedState] + + class Schema2(Schema1): + b: Annotated[int, InjectedState("bar")] + + @dec_tool(args_schema=Schema2) + def foo(a: str, b: int) -> float: + """return""" + return 0.0 + + assert _get_state_args(foo) == {"a": None, "b": "bar"} diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index a73083290..fda881b15 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8,6 +8,7 @@ import warnings from collections import Counter from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager +from dataclasses import replace from random import randrange from typing import ( Annotated, @@ -54,12 +55,17 @@ from langgraph.checkpoint.base import ( CheckpointTuple, ) from langgraph.checkpoint.memory import MemorySaver -from langgraph.constants import ERROR, PULL, PUSH +from langgraph.constants import ( + CONFIG_KEY_NODE_FINISHED, + ERROR, + FF_SEND_V2, + PULL, + PUSH, + START, +) from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt -from langgraph.graph import END, Graph -from langgraph.graph.graph import START -from langgraph.graph.message import MessageGraph, add_messages -from langgraph.graph.state import StateGraph +from langgraph.graph import END, Graph, GraphCommand, StateGraph +from langgraph.graph.message import MessageGraph, MessagesState, add_messages from langgraph.managed.shared_value import SharedValue from langgraph.prebuilt.chat_agent_executor import ( create_tool_calling_executor, @@ -74,7 +80,14 @@ from langgraph.pregel import ( from langgraph.pregel.retry import RetryPolicy from langgraph.store.base import BaseStore from langgraph.store.memory import InMemoryStore -from langgraph.types import Interrupt, PregelTask, Send, StreamWriter +from langgraph.types import ( + Command, + Interrupt, + PregelTask, + Send, + StreamWriter, + interrupt, +) from tests.any_str import AnyDict, AnyStr, AnyVersion, FloatBetween, UnsortedSequence from tests.conftest import ( ALL_CHECKPOINTERS_SYNC, @@ -748,7 +761,13 @@ def test_invoke_two_processes_in_out_interrupt( "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, ), @@ -768,6 +787,7 @@ def test_invoke_two_processes_in_out_interrupt( "source": "loop", "step": 5, "writes": {"one": None}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=history[2].config, @@ -788,6 +808,7 @@ def test_invoke_two_processes_in_out_interrupt( "source": "input", "step": 4, "writes": {"input": 3}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=history[3].config, @@ -808,6 +829,7 @@ def test_invoke_two_processes_in_out_interrupt( "source": "loop", "step": 3, "writes": {"one": None}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=history[4].config, @@ -828,6 +850,7 @@ def test_invoke_two_processes_in_out_interrupt( "source": "input", "step": 2, "writes": {"input": 20}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=history[5].config, @@ -843,7 +866,13 @@ def test_invoke_two_processes_in_out_interrupt( "checkpoint_id": AnyStr(), } }, - metadata={"parents": {}, "source": "loop", "step": 1, "writes": {"two": 4}}, + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": {"two": 4}, + "thread_id": "1", + }, created_at=AnyStr(), parent_config=history[6].config, ), @@ -863,6 +892,7 @@ def test_invoke_two_processes_in_out_interrupt( "source": "loop", "step": 0, "writes": {"one": None}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=history[7].config, @@ -883,6 +913,7 @@ def test_invoke_two_processes_in_out_interrupt( "source": "input", "step": -1, "writes": {"input": 2}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=None, @@ -949,6 +980,7 @@ def test_fork_always_re_runs_nodes( "source": "loop", "step": 5, "writes": {"add_one": 1}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=history[1].config, @@ -969,6 +1001,7 @@ def test_fork_always_re_runs_nodes( "source": "loop", "step": 4, "writes": {"add_one": 1}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=history[2].config, @@ -989,6 +1022,7 @@ def test_fork_always_re_runs_nodes( "source": "loop", "step": 3, "writes": {"add_one": 1}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=history[3].config, @@ -1009,6 +1043,7 @@ def test_fork_always_re_runs_nodes( "source": "loop", "step": 2, "writes": {"add_one": 1}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=history[4].config, @@ -1029,6 +1064,7 @@ def test_fork_always_re_runs_nodes( "source": "loop", "step": 1, "writes": {"add_one": 1}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=history[5].config, @@ -1044,7 +1080,13 @@ def test_fork_always_re_runs_nodes( "checkpoint_id": AnyStr(), } }, - metadata={"parents": {}, "source": "loop", "step": 0, "writes": None}, + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, created_at=AnyStr(), parent_config=history[6].config, ), @@ -1064,6 +1106,7 @@ def test_fork_always_re_runs_nodes( "source": "input", "step": -1, "writes": {"__start__": 1}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=None, @@ -1087,6 +1130,86 @@ def test_fork_always_re_runs_nodes( ] +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_run_from_checkpoint_id_retains_previous_writes( + request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + class MyState(TypedDict): + myval: Annotated[int, operator.add] + otherval: bool + + class Anode: + def __init__(self): + self.switch = False + + def __call__(self, state: MyState): + self.switch = not self.switch + return {"myval": 2 if self.switch else 1, "otherval": self.switch} + + builder = StateGraph(MyState) + thenode = Anode() # Fun. + builder.add_node("node_one", thenode) + builder.add_node("node_two", thenode) + builder.add_edge(START, "node_one") + + def _getedge(src: str): + swap = "node_one" if src == "node_two" else "node_two" + + def _edge(st: MyState) -> Literal["__end__", "node_one", "node_two"]: + if st["myval"] > 3: + return END + if st["otherval"]: + return swap + return src + + return _edge + + builder.add_conditional_edges("node_one", _getedge("node_one")) + builder.add_conditional_edges("node_two", _getedge("node_two")) + graph = builder.compile(checkpointer=checkpointer) + + thread_id = uuid.uuid4() + thread1 = {"configurable": {"thread_id": str(thread_id)}} + + result = graph.invoke({"myval": 1}, thread1) + assert result["myval"] == 4 + history = [c for c in graph.get_state_history(thread1)] + + assert len(history) == 4 + assert history[-1].values == {"myval": 0} + assert history[0].values == {"myval": 4, "otherval": False} + + second_run_config = { + **thread1, + "configurable": { + **thread1["configurable"], + "checkpoint_id": history[1].config["configurable"]["checkpoint_id"], + }, + } + second_result = graph.invoke(None, second_run_config) + assert second_result == {"myval": 5, "otherval": True} + + new_history = [ + c + for c in graph.get_state_history( + {"configurable": {"thread_id": str(thread_id), "checkpoint_ns": ""}} + ) + ] + + assert len(new_history) == len(history) + 1 + for original, new in zip(history, new_history[1:]): + assert original.values == new.values + assert original.next == new.next + assert original.metadata["step"] == new.metadata["step"] + + def _get_tasks(hist: list, start: int): + return [h.tasks for h in hist[start:]] + + assert _get_tasks(new_history, 1) == _get_tasks(history, 0) + + def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") @@ -1483,10 +1606,11 @@ def test_pending_writes_resume( assert two.calls == 2 # two attempts # latest checkpoint should be before nodes "one", "two" + # but we should have applied the write from "one" state = graph.get_state(thread1) assert state is not None - assert state.values == {"value": 1} - assert state.next == ("one", "two") + 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")'), @@ -1496,7 +1620,13 @@ def test_pending_writes_resume( "source": "loop", "step": 0, "writes": None, + "thread_id": "1", } + # get_state with checkpoint_id should not apply any pending writes + state = graph.get_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 = checkpointer.get_tuple(thread1) assert checkpoint is not None @@ -1584,6 +1714,7 @@ def test_pending_writes_resume( "step": 1, "source": "loop", "writes": {"one": {"value": 2}, "two": {"value": 3}}, + "thread_id": "1", }, parent_config={ "configurable": { @@ -1628,7 +1759,13 @@ def test_pending_writes_resume( "start:two": "__start__", }, }, - metadata={"parents": {}, "step": 0, "source": "loop", "writes": None}, + metadata={ + "parents": {}, + "step": 0, + "source": "loop", + "writes": None, + "thread_id": "1", + }, parent_config={ "configurable": { "thread_id": "1", @@ -1668,6 +1805,7 @@ def test_pending_writes_resume( "step": -1, "source": "input", "writes": {"__start__": {"value": 1}}, + "thread_id": "1", }, parent_config=None, pending_writes=UnsortedSequence( @@ -1704,6 +1842,1199 @@ def test_cond_edge_after_send() -> None: assert graph.invoke(["0"]) == ["0", "1", "2", "2", "3"] +def test_concurrent_emit_sends() -> None: + class Node: + def __init__(self, name: str): + self.name = name + setattr(self, "__name__", name) + + def __call__(self, state): + return ( + [self.name] + if isinstance(state, list) + else ["|".join((self.name, str(state)))] + ) + + def send_for_fun(state): + return [Send("2", 1), Send("2", 2), "3.1"] + + def send_for_profit(state): + return [Send("2", 3), Send("2", 4)] + + def route_to_three(state) -> Literal["3"]: + return "3" + + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node(Node("1")) + builder.add_node(Node("1.1")) + builder.add_node(Node("2")) + builder.add_node(Node("3")) + builder.add_node(Node("3.1")) + builder.add_edge(START, "1") + builder.add_edge(START, "1.1") + builder.add_conditional_edges("1", send_for_fun) + builder.add_conditional_edges("1.1", send_for_profit) + builder.add_conditional_edges("2", route_to_three) + graph = builder.compile() + assert graph.invoke(["0"]) == ( + [ + "0", + "1", + "1.1", + "2|1", + "2|2", + "2|3", + "2|4", + "3", + "3.1", + ] + if FF_SEND_V2 + else [ + "0", + "1", + "1.1", + "3.1", + "2|1", + "2|2", + "2|3", + "2|4", + "3", + ] + ) + + +def test_send_sequences() -> None: + class Node: + def __init__(self, name: str): + self.name = name + setattr(self, "__name__", name) + + def __call__(self, state): + update = ( + [self.name] + if isinstance(state, list) + else ["|".join((self.name, str(state)))] + ) + if isinstance(state, Command): + return replace(state, update=update) + else: + return update + + def send_for_fun(state): + return [ + Send("2", GraphCommand(send=Send("2", 3))), + Send("2", GraphCommand(send=Send("2", 4))), + "3.1", + ] + + def route_to_three(state) -> Literal["3"]: + return "3" + + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node(Node("1")) + builder.add_node(Node("2")) + builder.add_node(Node("3")) + builder.add_node(Node("3.1")) + builder.add_edge(START, "1") + builder.add_conditional_edges("1", send_for_fun) + builder.add_conditional_edges("2", route_to_three) + graph = builder.compile() + assert ( + graph.invoke(["0"]) + == [ + "0", + "1", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='2', arg=4))", + "2|3", + "2|4", + "3", + "3.1", + ] + if FF_SEND_V2 + else [ + "0", + "1", + "3.1", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='2', arg=4))", + "3", + "2|3", + "2|4", + "3", + ] + ) + + +@pytest.mark.repeat(20) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_send_dedupe_on_resume( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + if not FF_SEND_V2: + pytest.skip("Send deduplication is only available in Send V2") + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + class InterruptOnce: + ticks: int = 0 + + def __call__(self, state): + self.ticks += 1 + if self.ticks == 1: + raise NodeInterrupt("Bahh") + return ["|".join(("flaky", str(state)))] + + class Node: + def __init__(self, name: str): + self.name = name + self.ticks = 0 + setattr(self, "__name__", name) + + def __call__(self, state): + self.ticks += 1 + update = ( + [self.name] + if isinstance(state, list) + else ["|".join((self.name, str(state)))] + ) + if isinstance(state, GraphCommand): + return replace(state, update=update) + else: + return update + + def send_for_fun(state): + return [ + Send("2", GraphCommand(send=Send("2", 3))), + Send("2", GraphCommand(send=Send("flaky", 4))), + "3.1", + ] + + def route_to_three(state) -> Literal["3"]: + return "3" + + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node(Node("1")) + builder.add_node(Node("2")) + builder.add_node(Node("3")) + builder.add_node(Node("3.1")) + builder.add_node("flaky", InterruptOnce()) + builder.add_edge(START, "1") + builder.add_conditional_edges("1", send_for_fun) + builder.add_conditional_edges("2", route_to_three) + + graph = builder.compile(checkpointer=checkpointer) + thread1 = {"configurable": {"thread_id": "1"}} + assert graph.invoke(["0"], thread1, debug=1) == [ + "0", + "1", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='flaky', arg=4))", + "2|3", + ] + assert builder.nodes["2"].runnable.func.ticks == 3 + assert builder.nodes["flaky"].runnable.func.ticks == 1 + # check state + state = graph.get_state(thread1) + assert state.next == ("flaky",) + # check history + history = [c for c in graph.get_state_history(thread1)] + assert len(history) == 2 + # resume execution + assert graph.invoke(None, thread1, debug=1) == [ + "0", + "1", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='flaky', arg=4))", + "2|3", + "flaky|4", + "3", + "3.1", + ] + # node "2" doesn't get called again, as we recover writes saved before + assert builder.nodes["2"].runnable.func.ticks == 3 + # node "flaky" gets called again, as it was interrupted + assert builder.nodes["flaky"].runnable.func.ticks == 2 + # check state + state = graph.get_state(thread1) + assert state.next == () + # check history + history = [c for c in graph.get_state_history(thread1)] + assert ( + history[1] + == [ + StateSnapshot( + values=[ + "0", + "1", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='flaky', arg=4))", + "2|3", + "flaky|4", + "3", + "3.1", + ], + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"3": ["3"], "3.1": ["3.1"]}, + "thread_id": "1", + "step": 2, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=(), + ), + StateSnapshot( + values=[ + "0", + "1", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='flaky', arg=4))", + "2|3", + "flaky|4", + ], + next=("3", "3.1"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "1": ["1"], + "2": [ + ["2|Command(send=Send(node='2', arg=3))"], + ["2|Command(send=Send(node='flaky', arg=4))"], + ["2|3"], + ], + "flaky": ["flaky|4"], + }, + "thread_id": "1", + "step": 1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="3", + path=("__pregel_pull", "3"), + error=None, + interrupts=(), + state=None, + result=["3"], + ), + PregelTask( + id=AnyStr(), + name="3.1", + path=("__pregel_pull", "3.1"), + error=None, + interrupts=(), + state=None, + result=["3.1"], + ), + ), + ), + StateSnapshot( + values=["0"], + next=("1", "2", "2", "2", "flaky"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": None, + "thread_id": "1", + "step": 0, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="1", + path=("__pregel_pull", "1"), + error=None, + interrupts=(), + state=None, + result=["1"], + ), + PregelTask( + id=AnyStr(), + name="2", + path=( + "__pregel_push", + ("__pregel_pull", "1"), + 2, + AnyStr(), + ), + error=None, + interrupts=(), + state=None, + result=["2|Command(send=Send(node='2', arg=3))"], + ), + PregelTask( + id=AnyStr(), + name="2", + path=( + "__pregel_push", + ("__pregel_pull", "1"), + 3, + AnyStr(), + ), + error=None, + interrupts=(), + state=None, + result=["2|Command(send=Send(node='flaky', arg=4))"], + ), + PregelTask( + id=AnyStr(), + name="2", + path=( + "__pregel_push", + ( + "__pregel_push", + ("__pregel_pull", "1"), + 2, + AnyStr(), + ), + 2, + AnyStr(), + ), + error=None, + interrupts=(), + state=None, + result=["2|3"], + ), + PregelTask( + id=AnyStr(), + name="flaky", + path=( + "__pregel_push", + ( + "__pregel_push", + ("__pregel_pull", "1"), + 3, + AnyStr(), + ), + 2, + AnyStr(), + ), + error=None, + interrupts=(Interrupt(value="Bahh", when="during"),), + state=None, + result=["flaky|4"], + ), + ), + ), + StateSnapshot( + values=[], + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"__start__": ["0"]}, + "thread_id": "1", + "step": -1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=("__pregel_pull", "__start__"), + error=None, + interrupts=(), + state=None, + result=["0"], + ), + ), + ), + ][1] + ) + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_send_react_interrupt( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage + + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + ai_message = AIMessage( + "", + id="ai1", + tool_calls=[ToolCall(name="foo", args={"hi": [1, 2, 3]}, id=AnyStr())], + ) + + def agent(state): + return {"messages": ai_message} + + def route(state): + if isinstance(state["messages"][-1], AIMessage): + return [ + Send(call["name"], call) for call in state["messages"][-1].tool_calls + ] + + foo_called = 0 + + def foo(call: ToolCall): + nonlocal foo_called + foo_called += 1 + return {"messages": ToolMessage(str(call["args"]), tool_call_id=call["id"])} + + builder = StateGraph(MessagesState) + builder.add_node(agent) + builder.add_node(foo) + builder.add_edge(START, "agent") + builder.add_conditional_edges("agent", route) + graph = builder.compile() + + assert graph.invoke({"messages": [HumanMessage("hello")]}) == { + "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=checkpointer, interrupt_before=["foo"]) + thread1 = {"configurable": {"thread_id": "1"}} + assert graph.invoke({"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 graph.invoke(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 graph.invoke({"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 + + if not FF_SEND_V2: + return + + # get state should show the pending task + state = graph.get_state(thread1) + assert state == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + }, + next=("foo",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 0, + "source": "loop", + "writes": None, + "parents": {}, + "thread_id": "2", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + content="", + additional_kwargs={}, + response_metadata={}, + id="ai1", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ) + }, + ), + PregelTask( + id=AnyStr(), + name="foo", + path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()), + error=None, + interrupts=(), + state=None, + result=None, + ), + ), + ) + + # remove the tool call, clearing the pending task + graph.update_state( + thread1, {"messages": AIMessage("Bye now", id=ai_message.id, tool_calls=[])} + ) + + # tool call no longer in pending tasks + assert graph.get_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": 1, + "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=(), + ) + + # tool call not executed + assert graph.invoke(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 graph.invoke({"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 = graph.get_state(thread1) + assert state == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + }, + next=("foo",), + config={ + "configurable": { + "thread_id": "3", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 0, + "source": "loop", + "writes": None, + "parents": {}, + "thread_id": "3", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "3", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + "", + id="ai1", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ) + }, + ), + PregelTask( + id=AnyStr(), + name="foo", + path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()), + error=None, + interrupts=(), + state=None, + result=None, + ), + ), + ) + + # replace the tool call, should clear previous send, create new one + graph.update_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 graph.get_state(thread1) == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [4, 5, 6]}, + "id": "tool1", + "type": "tool_call", + } + ], + ), + ] + }, + next=("foo",), + config={ + "configurable": { + "thread_id": "3", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 1, + "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, AnyStr()), + error=None, + interrupts=(), + state=None, + result=None, + ), + ), + ) + + # prev tool call not executed, new tool call is + assert graph.invoke(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 + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_send_react_interrupt_control( + request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion +) -> None: + from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage + + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + ai_message = AIMessage( + "", + id="ai1", + tool_calls=[ToolCall(name="foo", args={"hi": [1, 2, 3]}, id=AnyStr())], + ) + + def agent(state) -> GraphCommand[Literal["foo"]]: + return GraphCommand( + update={"messages": ai_message}, + send=[Send(call["name"], call) for call in ai_message.tool_calls], + ) + + foo_called = 0 + + def foo(call: ToolCall): + nonlocal foo_called + foo_called += 1 + return {"messages": ToolMessage(str(call["args"]), tool_call_id=call["id"])} + + builder = StateGraph(MessagesState) + builder.add_node(agent) + builder.add_node(foo) + builder.add_edge(START, "agent") + graph = builder.compile() + assert graph.get_graph().draw_mermaid() == snapshot + + assert graph.invoke({"messages": [HumanMessage("hello")]}) == { + "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=checkpointer, interrupt_before=["foo"]) + thread1 = {"configurable": {"thread_id": "1"}} + assert graph.invoke({"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 graph.invoke(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 + + if not FF_SEND_V2: + return + + # interrupt-update-resume flow + foo_called = 0 + graph = builder.compile(checkpointer=checkpointer, interrupt_before=["foo"]) + thread1 = {"configurable": {"thread_id": "2"}} + assert graph.invoke({"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 = graph.get_state(thread1) + assert state == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + }, + next=("foo",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 0, + "source": "loop", + "writes": None, + "parents": {}, + "thread_id": "2", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + content="", + additional_kwargs={}, + response_metadata={}, + id="ai1", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ) + }, + ), + PregelTask( + id=AnyStr(), + name="foo", + path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()), + error=None, + interrupts=(), + state=None, + result=None, + ), + ), + ) + + # remove the tool call, clearing the pending task + graph.update_state( + thread1, {"messages": AIMessage("Bye now", id=ai_message.id, tool_calls=[])} + ) + + # tool call no longer in pending tasks + assert graph.get_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": 1, + "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=(), + ) + + # tool call not executed + assert graph.invoke(None, thread1) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage(content="Bye now"), + ] + } + assert foo_called == 0 + + # interrupt-update-resume flow, creating new Send in update call + + # TODO add here test with invoke(Command()) + + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_invoke_checkpoint_three( mocker: MockerFixture, request: pytest.FixtureRequest, checkpointer_name: str @@ -2117,7 +3448,9 @@ def test_conditional_graph( workflow.add_node("agent", agent) workflow.add_node( - "tools", execute_tools, metadata={"parents": {}, "version": 2, "variant": "b"} + "tools", + execute_tools, + metadata={"parents": {}, "version": 2, "variant": "b"}, ) workflow.set_entry_point("agent") @@ -2311,6 +3644,7 @@ def test_conditional_graph( } }, }, + "thread_id": "1", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -2362,6 +3696,7 @@ def test_conditional_graph( "input": "what is weather in sf", }, }, + "thread_id": "1", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -2482,6 +3817,7 @@ def test_conditional_graph( ), } }, + "thread_id": "1", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -2537,6 +3873,7 @@ def test_conditional_graph( } } }, + "thread_id": "2", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -2582,6 +3919,7 @@ def test_conditional_graph( "input": "what is weather in sf", } }, + "thread_id": "2", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -2702,6 +4040,7 @@ def test_conditional_graph( ), } }, + "thread_id": "2", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -2757,6 +4096,7 @@ def test_conditional_graph( } } }, + "thread_id": "3", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -3230,6 +4570,7 @@ def test_conditional_state_graph( ), } }, + "thread_id": "1", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -3272,6 +4613,7 @@ def test_conditional_state_graph( ) }, }, + "thread_id": "1", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -3348,6 +4690,7 @@ def test_conditional_state_graph( ) } }, + "thread_id": "1", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -3399,6 +4742,7 @@ def test_conditional_state_graph( ), } }, + "thread_id": "2", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -3440,6 +4784,7 @@ def test_conditional_state_graph( ) } }, + "thread_id": "2", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -3514,6 +4859,7 @@ def test_conditional_state_graph( ) } }, + "thread_id": "2", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -3541,7 +4887,13 @@ def test_conditional_state_graph( next=("agent",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], - metadata={"parents": {}, "source": "loop", "step": 0, "writes": None}, + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "3", + }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -3580,6 +4932,7 @@ def test_conditional_state_graph( ), } }, + "thread_id": "3", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -3640,6 +4993,7 @@ def test_conditional_state_graph( ], } }, + "thread_id": "3", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -3702,6 +5056,7 @@ def test_conditional_state_graph( ), } }, + "thread_id": "4", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -3762,6 +5117,7 @@ def test_conditional_state_graph( ], } }, + "thread_id": "4", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -4078,18 +5434,6 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: stream_mode="messages", ) ] == [ - ( - _AnyIdHumanMessage( - content="what is weather in sf", - ), - { - "langgraph_step": 0, - "langgraph_node": "__start__", - "langgraph_triggers": ["__start__"], - "langgraph_path": ("__pregel_pull", "__start__"), - "langgraph_checkpoint_ns": AnyStr("__start__:"), - }, - ), ( _AnyIdAIMessageChunk( content="", @@ -4115,7 +5459,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: "langgraph_step": 1, "langgraph_node": "agent", "langgraph_triggers": ["start:agent"], - "langgraph_path": ("__pregel_pull", "agent"), + "langgraph_path": (PULL, "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), "ls_provider": "fakechatmodel", @@ -4132,7 +5476,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: "langgraph_step": 2, "langgraph_node": "tools", "langgraph_triggers": ["branch:agent:should_continue:tools"], - "langgraph_path": ("__pregel_pull", "tools"), + "langgraph_path": (PULL, "tools"), "langgraph_checkpoint_ns": AnyStr("tools:"), }, ), @@ -4174,7 +5518,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: "langgraph_step": 3, "langgraph_node": "agent", "langgraph_triggers": ["tools"], - "langgraph_path": ("__pregel_pull", "agent"), + "langgraph_path": (PULL, "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), "ls_provider": "fakechatmodel", @@ -4191,7 +5535,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: "langgraph_step": 4, "langgraph_node": "tools", "langgraph_triggers": ["branch:agent:should_continue:tools"], - "langgraph_path": ("__pregel_pull", "tools"), + "langgraph_path": (PULL, "tools"), "langgraph_checkpoint_ns": AnyStr("tools:"), }, ), @@ -4205,7 +5549,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: "langgraph_step": 4, "langgraph_node": "tools", "langgraph_triggers": ["branch:agent:should_continue:tools"], - "langgraph_path": ("__pregel_pull", "tools"), + "langgraph_path": (PULL, "tools"), "langgraph_checkpoint_ns": AnyStr("tools:"), }, ), @@ -4217,7 +5561,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: "langgraph_step": 5, "langgraph_node": "agent", "langgraph_triggers": ["tools"], - "langgraph_path": ("__pregel_pull", "agent"), + "langgraph_path": (PULL, "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), "ls_provider": "fakechatmodel", @@ -4620,6 +5964,8 @@ def test_state_graph_packets( {"agent": {"messages": AIMessage(content="answer", id="ai3")}}, ] + # interrupt after agent + app_w_interrupt = workflow.compile( checkpointer=checkpointer, interrupt_after=["agent"], @@ -4650,6 +5996,9 @@ def test_state_graph_packets( {"__interrupt__": ()}, ] + if not FF_SEND_V2: + return + assert app_w_interrupt.get_state(config) == StateSnapshot( values={ "messages": [ @@ -4667,29 +6016,44 @@ def test_state_graph_packets( ), ] }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),), + tasks=( + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + content="", + additional_kwargs={}, + response_metadata={}, + id="ai1", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "query"}, + "id": "tool_call123", + "type": "tool_call", + } + ], + ) + }, + ), + PregelTask( + AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) + ), + ), next=("tools",), config=(app_w_interrupt.checkpointer.get_tuple(config)).config, created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], metadata={ "parents": {}, "source": "loop", - "step": 1, - "writes": { - "agent": { - "messages": AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ) - } - }, + "step": 0, + "writes": None, + "thread_id": "1", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -4719,14 +6083,14 @@ def test_state_graph_packets( ), ] }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),), + tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0, AnyStr())),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], metadata={ "parents": {}, "source": "update", - "step": 2, + "step": 1, "writes": { "agent": { "messages": AIMessage( @@ -4743,6 +6107,7 @@ def test_state_graph_packets( "something_extra": "hi there", } }, + "thread_id": "1", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -4819,8 +6184,40 @@ def test_state_graph_packets( ] }, tasks=( - PregelTask(AnyStr(), "tools", (PUSH, 0)), - PregelTask(AnyStr(), "tools", (PUSH, 1)), + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + "", + id="ai2", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "another", "idx": 0}, + "id": "tool_call234", + "type": "tool_call", + }, + { + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + "id": "tool_call567", + "type": "tool_call", + }, + ], + ) + }, + ), + PregelTask( + AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) + ), + PregelTask( + AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3, AnyStr()) + ), ), next=("tools", "tools"), config=app_w_interrupt.checkpointer.get_tuple(config).config, @@ -4828,27 +6225,17 @@ def test_state_graph_packets( metadata={ "parents": {}, "source": "loop", - "step": 4, + "step": 2, "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}, - }, - ], - ) + "tools": { + "messages": _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), }, }, + "thread_id": "1", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -4892,13 +6279,337 @@ def test_state_graph_packets( metadata={ "parents": {}, "source": "update", - "step": 5, + "step": 3, "writes": { "agent": { "messages": AIMessage(content="answer", id="ai2"), "something_extra": "hi there", } }, + "thread_id": "1", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + # interrupt before tools + + app_w_interrupt = workflow.compile( + checkpointer=checkpointer, + interrupt_before=["tools"], + ) + config = {"configurable": {"thread_id": "2"}} + model.i = 0 + + assert [ + c + for c in app_w_interrupt.stream( + {"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__": ()}, + ] + + assert app_w_interrupt.get_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( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + "", + id="ai1", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "query"}, + "id": "tool_call123", + "type": "tool_call", + } + ], + ) + }, + ), + PregelTask( + AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) + ), + ), + next=("tools",), + config=(app_w_interrupt.checkpointer.get_tuple(config)).config, + created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "2", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + # modify ai message + last_message = (app_w_interrupt.get_state(config)).values["messages"][-1] + last_message.tool_calls[0]["args"]["query"] = "a different query" + app_w_interrupt.update_state( + config, {"messages": last_message, "something_extra": "hi there"} + ) + + # message was replaced instead of appended + assert app_w_interrupt.get_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, AnyStr())),), + next=("tools",), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + "something_extra": "hi there", + } + }, + "thread_id": "2", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + assert [c for c in app_w_interrupt.stream(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__": ()}, + ] + + assert app_w_interrupt.get_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( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + "", + id="ai2", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "another", "idx": 0}, + "id": "tool_call234", + "type": "tool_call", + }, + { + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + "id": "tool_call567", + "type": "tool_call", + }, + ], + ) + }, + ), + PregelTask( + AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) + ), + PregelTask( + AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3, AnyStr()) + ), + ), + next=("tools", "tools"), + config=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 2, + "writes": { + "tools": { + "messages": _AnyIdToolMessage( + content="result for a different query", + name="search_api", + tool_call_id="tool_call123", + ), + }, + }, + "thread_id": "2", + }, + parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, + ) + + app_w_interrupt.update_state( + config, + { + "messages": AIMessage(content="answer", id="ai2"), + "something_extra": "hi there", + }, + ) + + # replaces message even if object identity is different, as long as id is the same + assert app_w_interrupt.get_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=app_w_interrupt.checkpointer.get_tuple(config).config, + created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 3, + "writes": { + "agent": { + "messages": AIMessage(content="answer", id="ai2"), + "something_extra": "hi there", + } + }, + "thread_id": "2", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -5188,6 +6899,7 @@ def test_message_graph( id="ai1", ) }, + "thread_id": "1", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -5234,6 +6946,7 @@ def test_message_graph( id="ai1", ) }, + "thread_id": "1", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -5316,6 +7029,7 @@ def test_message_graph( id="ai2", ) }, + "thread_id": "1", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -5356,6 +7070,7 @@ def test_message_graph( "source": "update", "step": 5, "writes": {"agent": AIMessage(content="answer", id="ai2")}, + "thread_id": "1", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -5420,6 +7135,7 @@ def test_message_graph( id="ai1", ) }, + "thread_id": "2", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -5466,6 +7182,7 @@ def test_message_graph( id="ai1", ) }, + "thread_id": "2", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -5548,6 +7265,7 @@ def test_message_graph( id="ai2", ) }, + "thread_id": "2", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -5589,6 +7307,7 @@ def test_message_graph( "source": "update", "step": 5, "writes": {"agent": AIMessage(content="answer", id="ai2")}, + "thread_id": "2", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -5630,6 +7349,7 @@ def test_message_graph( "source": "update", "step": 6, "writes": {"tools": UnsortedSequence("ai", "an extra message")}, + "thread_id": "2", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -5918,6 +7638,7 @@ def test_root_graph( id="ai1", ) }, + "thread_id": "1", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -5964,6 +7685,7 @@ def test_root_graph( id="ai1", ) }, + "thread_id": "1", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -6047,6 +7769,7 @@ def test_root_graph( id="ai2", ) }, + "thread_id": "1", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -6088,6 +7811,7 @@ def test_root_graph( "source": "update", "step": 5, "writes": {"agent": AIMessage(content="answer", id="ai2")}, + "thread_id": "1", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -6152,6 +7876,7 @@ def test_root_graph( id="ai1", ) }, + "thread_id": "2", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -6198,6 +7923,7 @@ def test_root_graph( id="ai1", ) }, + "thread_id": "2", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -6281,6 +8007,7 @@ def test_root_graph( id="ai2", ) }, + "thread_id": "2", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -6321,6 +8048,7 @@ def test_root_graph( "source": "update", "step": 5, "writes": {"agent": AIMessage(content="answer", id="ai2")}, + "thread_id": "2", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -6362,6 +8090,7 @@ def test_root_graph( "source": "update", "step": 6, "writes": {"tools": UnsortedSequence("ai", "an extra message")}, + "thread_id": "2", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -6434,6 +8163,7 @@ def test_root_graph( "source": "update", "step": 6, "writes": {"tools": UnsortedSequence("ai", "an extra message")}, + "thread_id": "2", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -6470,7 +8200,7 @@ def test_root_graph( ), AIMessage(content="answer", id="ai2"), AIMessage( - content="an extra message", id="00000000-0000-4000-8000-000000000091" + content="an extra message", id="00000000-0000-4000-8000-000000000092" ), HumanMessage(content="what is weather in la"), ], @@ -6718,8 +8448,10 @@ def test_dynamic_interrupt( nonlocal tool_two_node_count tool_two_node_count += 1 if s["market"] == "DE": - raise NodeInterrupt("Just because...") - return {"my_key": " all good"} + answer = interrupt("Just because...") + else: + answer = " all good" + return {"my_key": answer} tool_two_graph = StateGraph(State) tool_two_graph.add_node("tool_two", tool_two_node, retry=RetryPolicy()) @@ -6751,6 +8483,28 @@ def test_dynamic_interrupt( with pytest.raises(ValueError, match="thread_id"): tool_two.invoke({"my_key": "value", "market": "DE"}) + # flow: interrupt -> resume with answer + thread2 = {"configurable": {"thread_id": "2"}} + # stop when about to enter node + assert [ + c for c in tool_two.stream({"my_key": "value ⛰️", "market": "DE"}, thread2) + ] == [ + { + "__interrupt__": ( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:")], + ), + ) + }, + ] + # resume with answer + assert [c for c in tool_two.stream(Command(resume=" my answer"), thread2)] == [ + {"tool_two": {"my_key": " my answer"}}, + ] + + # flow: interrupt -> clear tasks thread1 = {"configurable": {"thread_id": "1"}} # stop when about to enter node assert tool_two.invoke({"my_key": "value ⛰️", "market": "DE"}, thread1) == { @@ -6763,12 +8517,14 @@ def test_dynamic_interrupt( "source": "loop", "step": 0, "writes": None, + "thread_id": "1", }, { "parents": {}, "source": "input", "step": -1, "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, + "thread_id": "1", }, ] assert tool_two.get_state(thread1) == StateSnapshot( @@ -6779,12 +8535,203 @@ def test_dynamic_interrupt( AnyStr(), "tool_two", (PULL, "tool_two"), - interrupts=(Interrupt("Just because..."),), + interrupts=( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:")], + ), + ), ), ), config=tool_two.checkpointer.get_tuple(thread1).config, created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={"parents": {}, "source": "loop", "step": 0, "writes": None}, + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + # clear the interrupt and next tasks + tool_two.update_state(thread1, None, as_node=END) + # interrupt and next tasks are cleared + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️", "market": "DE"}, + next=(), + tasks=(), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": {}, + "thread_id": "1", + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + + +@pytest.mark.skipif(not FF_SEND_V2, reason="send v2 is not enabled") +@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}") + + class State(TypedDict): + my_key: Annotated[str, operator.add] + market: str + + def tool_one(s: State) -> State: + return {"my_key": " one"} + + tool_two_node_count = 0 + + def tool_two_node(s: State) -> State: + nonlocal tool_two_node_count + tool_two_node_count += 1 + if s["market"] == "DE": + answer = interrupt("Just because...") + else: + answer = " all good" + return {"my_key": answer} + + def start(state: State) -> list[Union[Send, str]]: + return ["tool_two", Send("tool_one", state)] + + tool_two_graph = StateGraph(State) + tool_two_graph.add_node("tool_two", tool_two_node, retry=RetryPolicy()) + tool_two_graph.add_node("tool_one", tool_one) + tool_two_graph.set_conditional_entry_point(start) + tool_two = tool_two_graph.compile() + + tracer = FakeTracer() + assert tool_two.invoke( + {"my_key": "value", "market": "DE"}, {"callbacks": [tracer]} + ) == { + "my_key": "value one", + "market": "DE", + } + assert tool_two_node_count == 1, "interrupts aren't retried" + assert len(tracer.runs) == 1 + run = tracer.runs[0] + assert run.end_time is not None + assert run.error is None + assert run.outputs == {"market": "DE", "my_key": "value one"} + + assert tool_two.invoke({"my_key": "value", "market": "US"}) == { + "my_key": "value one all good", + "market": "US", + } + + tool_two = tool_two_graph.compile(checkpointer=checkpointer) + + # missing thread_id + with pytest.raises(ValueError, match="thread_id"): + tool_two.invoke({"my_key": "value", "market": "DE"}) + + # flow: interrupt -> resume with answer + thread2 = {"configurable": {"thread_id": "2"}} + # stop when about to enter node + assert [ + c for c in tool_two.stream({"my_key": "value ⛰️", "market": "DE"}, thread2) + ] == [ + { + "tool_one": {"my_key": " one"}, + }, + { + "__interrupt__": ( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:")], + ), + ) + }, + ] + # resume with answer + assert [c for c in tool_two.stream(Command(resume=" my answer"), thread2)] == [ + {"tool_two": {"my_key": " my answer"}}, + ] + + # flow: interrupt -> clear tasks + thread1 = {"configurable": {"thread_id": "1"}} + # stop when about to enter node + assert tool_two.invoke({"my_key": "value ⛰️", "market": "DE"}, thread1) == { + "my_key": "value ⛰️ one", + "market": "DE", + } + assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": {"tool_one": {"my_key": " one"}}, + "thread_id": "1", + }, + { + "parents": {}, + "source": "input", + "step": -1, + "writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}}, + "thread_id": "1", + }, + ] + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️ one", "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=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": {"tool_one": {"my_key": " one"}}, + "thread_id": "1", + }, + parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, + ) + # clear the interrupt and next tasks + tool_two.update_state(thread1, None) + # interrupt is cleared, next task is kept + assert tool_two.get_state(thread1) == StateSnapshot( + values={"my_key": "value ⛰️ one", "market": "DE"}, + next=("tool_two",), + tasks=( + PregelTask( + AnyStr(), + "tool_two", + (PULL, "tool_two"), + interrupts=(), + ), + ), + config=tool_two.checkpointer.get_tuple(thread1).config, + created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "writes": {}, + "thread_id": "1", + }, parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, ) @@ -6861,12 +8808,16 @@ def test_start_branch_then( "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 tool_two.get_state(thread1) == StateSnapshot( @@ -6875,7 +8826,14 @@ def test_start_branch_then( next=("tool_two_slow",), config=tool_two.checkpointer.get_tuple(thread1).config, created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], - metadata={"parents": {}, "source": "loop", "step": 0, "writes": None}, + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "assistant_id": "a", + "thread_id": "1", + }, parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, ) # resume, for same result as above @@ -6894,6 +8852,8 @@ def test_start_branch_then( "source": "loop", "step": 1, "writes": {"tool_two_slow": {"my_key": " slow"}}, + "assistant_id": "a", + "thread_id": "1", }, parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, ) @@ -6910,7 +8870,14 @@ def test_start_branch_then( next=("tool_two_fast",), config=tool_two.checkpointer.get_tuple(thread2).config, created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], - metadata={"parents": {}, "source": "loop", "step": 0, "writes": None}, + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "assistant_id": "a", + "thread_id": "2", + }, parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, ) # resume, for same result as above @@ -6929,6 +8896,8 @@ def test_start_branch_then( "source": "loop", "step": 1, "writes": {"tool_two_fast": {"my_key": " fast"}}, + "assistant_id": "a", + "thread_id": "2", }, parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, ) @@ -6945,7 +8914,14 @@ def test_start_branch_then( next=("tool_two_fast",), config=tool_two.checkpointer.get_tuple(thread3).config, created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], - metadata={"parents": {}, "source": "loop", "step": 0, "writes": None}, + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "assistant_id": "b", + "thread_id": "3", + }, parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, ) # update state @@ -6961,6 +8937,8 @@ def test_start_branch_then( "source": "update", "step": 1, "writes": {START: {"my_key": "key"}}, + "assistant_id": "b", + "thread_id": "3", }, parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, ) @@ -6980,6 +8958,8 @@ def test_start_branch_then( "source": "loop", "step": 2, "writes": {"tool_two_fast": {"my_key": " fast"}}, + "assistant_id": "b", + "thread_id": "3", }, parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, ) @@ -7053,6 +9033,7 @@ def test_branch_then( "source": "input", "step": -1, "writes": {"__start__": {"my_key": "value", "market": "DE"}}, + "thread_id": "10", }, "parent_config": None, "next": ["__start__"], @@ -7091,6 +9072,7 @@ def test_branch_then( "source": "loop", "step": 0, "writes": None, + "thread_id": "10", }, "parent_config": { "tags": [], @@ -7157,6 +9139,7 @@ def test_branch_then( "source": "loop", "step": 1, "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "10", }, "parent_config": { "tags": [], @@ -7228,6 +9211,7 @@ def test_branch_then( "source": "loop", "step": 2, "writes": {"tool_two_slow": {"my_key": " slow"}}, + "thread_id": "10", }, "parent_config": { "tags": [], @@ -7294,6 +9278,7 @@ def test_branch_then( "source": "loop", "step": 3, "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "10", }, "parent_config": { "tags": [], @@ -7337,6 +9322,7 @@ def test_branch_then( "source": "loop", "step": 1, "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "1", }, parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, ) @@ -7356,6 +9342,7 @@ def test_branch_then( "source": "loop", "step": 3, "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "1", }, parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, ) @@ -7377,6 +9364,7 @@ def test_branch_then( "source": "loop", "step": 1, "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "2", }, parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, ) @@ -7396,6 +9384,7 @@ def test_branch_then( "source": "loop", "step": 3, "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "2", }, parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, ) @@ -7425,6 +9414,7 @@ def test_branch_then( "source": "loop", "step": 2, "writes": {"tool_two_slow": {"my_key": " slow"}}, + "thread_id": "11", }, parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, ) @@ -7445,6 +9435,7 @@ def test_branch_then( "source": "update", "step": 3, "writes": {"tool_two_slow": {"my_key": "er"}}, + "thread_id": "11", }, parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, ) @@ -7474,6 +9465,7 @@ def test_branch_then( "source": "loop", "step": 1, "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "21", }, parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, ) @@ -7493,6 +9485,7 @@ def test_branch_then( "source": "loop", "step": 3, "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "21", }, parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config, ) @@ -7514,6 +9507,7 @@ def test_branch_then( "source": "loop", "step": 1, "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "22", }, parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, ) @@ -7533,6 +9527,7 @@ def test_branch_then( "source": "loop", "step": 3, "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "22", }, parent_config=[*tool_two.checkpointer.list(thread2, limit=2)][-1].config, ) @@ -7552,6 +9547,7 @@ def test_branch_then( "source": "update", "step": 0, "writes": {START: {"my_key": "key", "market": "DE"}}, + "thread_id": "23", }, parent_config=None, ) @@ -7572,6 +9568,7 @@ def test_branch_then( "source": "loop", "step": 1, "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "23", }, parent_config=uconfig, ) @@ -7591,6 +9588,7 @@ def test_branch_then( "source": "loop", "step": 3, "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "23", }, parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, ) @@ -7718,6 +9716,7 @@ def test_in_one_fan_out_state_graph_waiting_edge( "source": "update", "step": 4, "writes": {"retriever_one": {"docs": ["doc5"]}}, + "thread_id": "2", }, parent_config=[*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config, ) @@ -8701,14 +10700,14 @@ def test_stream_subgraphs_during_execution( ), (FloatBetween(0.2, 0.3), ((), {"outer_1": {"my_key": " and parallel"}})), ( - FloatBetween(0.5, 0.6), + FloatBetween(0.5, 0.8), ( (AnyStr("inner:"),), {"inner_2": {"my_key": " and there", "my_other_key": "got here"}}, ), ), - (FloatBetween(0.5, 0.6), ((), {"inner": {"my_key": "got here and there"}})), - (FloatBetween(0.5, 0.6), ((), {"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"}})), ] @@ -8796,7 +10795,7 @@ def test_nested_graph_interrupts_parallel( # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} assert app.invoke({"my_key": ""}, config, debug=True) == { - "my_key": "", + "my_key": " and parallel", } assert app.invoke(None, config, debug=True) == { @@ -8824,6 +10823,7 @@ def test_nested_graph_interrupts_parallel( config = {"configurable": {"thread_id": "3"}} assert [*app.stream({"my_key": ""}, config, stream_mode="values")] == [ {"my_key": ""}, + {"my_key": " and parallel"}, ] assert [*app.stream(None, config, stream_mode="values")] == [ {"my_key": ""}, @@ -8840,6 +10840,7 @@ def test_nested_graph_interrupts_parallel( # while we're waiting for the node w/ interrupt inside to finish assert [*app.stream(None, config, stream_mode="values")] == [ {"my_key": ""}, + {"my_key": " and parallel"}, ] assert [*app.stream(None, config, stream_mode="values")] == [ {"my_key": ""}, @@ -8851,7 +10852,8 @@ def test_nested_graph_interrupts_parallel( app = graph.compile(checkpointer=checkpointer, interrupt_after=["outer_1"]) config = {"configurable": {"thread_id": "5"}} assert [*app.stream({"my_key": ""}, config, stream_mode="values")] == [ - {"my_key": ""} + {"my_key": ""}, + {"my_key": " and parallel"}, ] assert [*app.stream(None, config, stream_mode="values")] == [ {"my_key": ""}, @@ -8929,15 +10931,27 @@ def test_doubly_nested_graph_interrupts( } # test stream updates w/ nested interrupt - config = {"configurable": {"thread_id": "2"}} + nodes: list[str] = [] + config = { + "configurable": {"thread_id": "2", CONFIG_KEY_NODE_FINISHED: nodes.append} + } assert [*app.stream({"my_key": "my value"}, config)] == [ {"parent_1": {"my_key": "hi my value"}}, {"__interrupt__": ()}, ] + assert nodes == ["parent_1", "grandchild_1"] assert [*app.stream(None, config)] == [ {"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"}} @@ -9032,6 +11046,7 @@ def test_nested_graph_state( "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, "step": 1, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -9085,6 +11100,13 @@ def test_nested_graph_state( } }, "step": 1, + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "langgraph_node": "inner", + "langgraph_path": [PULL, "inner"], + "langgraph_step": 2, + "langgraph_triggers": ["outer_1"], + "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), parent_config={ @@ -9113,6 +11135,7 @@ def test_nested_graph_state( "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, "step": 1, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -9154,6 +11177,7 @@ def test_nested_graph_state( "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, "step": 1, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -9182,7 +11206,13 @@ def test_nested_graph_state( "checkpoint_id": AnyStr(), } }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, + metadata={ + "parents": {}, + "source": "loop", + "writes": None, + "step": 0, + "thread_id": "1", + }, created_at=AnyStr(), parent_config={ "configurable": { @@ -9215,6 +11245,7 @@ def test_nested_graph_state( "source": "input", "writes": {"__start__": {"my_key": "my value"}}, "step": -1, + "thread_id": "1", }, created_at=AnyStr(), parent_config=None, @@ -9246,6 +11277,13 @@ def test_nested_graph_state( }, "step": 1, "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "langgraph_node": "inner", + "langgraph_path": [PULL, "inner"], + "langgraph_step": 2, + "langgraph_triggers": ["outer_1"], + "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), parent_config={ @@ -9278,6 +11316,13 @@ def test_nested_graph_state( "writes": None, "step": 0, "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "langgraph_node": "inner", + "langgraph_path": [PULL, "inner"], + "langgraph_step": 2, + "langgraph_triggers": ["outer_1"], + "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), parent_config={ @@ -9320,6 +11365,13 @@ def test_nested_graph_state( "writes": {"__start__": {"my_key": "hi my value"}}, "step": -1, "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "langgraph_node": "inner", + "langgraph_path": [PULL, "inner"], + "langgraph_step": 2, + "langgraph_triggers": ["outer_1"], + "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), parent_config=None, @@ -9355,6 +11407,7 @@ def test_nested_graph_state( "outer_2": {"my_key": "hi my value here and there and back again"} }, "step": 3, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -9386,6 +11439,7 @@ def test_nested_graph_state( "outer_2": {"my_key": "hi my value here and there and back again"} }, "step": 3, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -9419,6 +11473,7 @@ def test_nested_graph_state( "source": "loop", "writes": {"inner": {"my_key": "hi my value here and there"}}, "step": 2, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -9455,6 +11510,7 @@ def test_nested_graph_state( "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, "step": 1, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -9483,7 +11539,13 @@ def test_nested_graph_state( "checkpoint_id": AnyStr(), } }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, + metadata={ + "parents": {}, + "source": "loop", + "writes": None, + "step": 0, + "thread_id": "1", + }, created_at=AnyStr(), parent_config={ "configurable": { @@ -9516,6 +11578,7 @@ def test_nested_graph_state( "source": "input", "writes": {"__start__": {"my_key": "my value"}}, "step": -1, + "thread_id": "1", }, created_at=AnyStr(), parent_config=None, @@ -9622,6 +11685,7 @@ def test_doubly_nested_graph_state( "source": "loop", "writes": {"parent_1": {"my_key": "hi my value"}}, "step": 1, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -9663,6 +11727,7 @@ def test_doubly_nested_graph_state( "source": "loop", "writes": None, "step": 0, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -9709,6 +11774,13 @@ def test_doubly_nested_graph_state( "source": "loop", "writes": {"grandchild_1": {"my_key": "hi my value here"}}, "step": 1, + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [PULL, AnyStr("child_1")], + "langgraph_step": 1, + "langgraph_triggers": [AnyStr("start:child_1")], }, created_at=AnyStr(), parent_config={ @@ -9779,6 +11851,16 @@ def test_doubly_nested_graph_state( "grandchild_1": {"my_key": "hi my value here"} }, "step": 1, + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [AnyStr("start:child_1")], }, created_at=AnyStr(), parent_config={ @@ -9816,6 +11898,13 @@ def test_doubly_nested_graph_state( "source": "loop", "writes": None, "step": 0, + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child", + "langgraph_path": [PULL, AnyStr("child")], + "langgraph_step": 2, + "langgraph_triggers": [AnyStr("parent_1")], + "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), parent_config={ @@ -9844,6 +11933,7 @@ def test_doubly_nested_graph_state( "source": "loop", "writes": {"parent_1": {"my_key": "hi my value"}}, "step": 1, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -9854,7 +11944,7 @@ def test_doubly_nested_graph_state( } }, ) - # resume + # # resume assert [c for c in app.stream(None, config, subgraphs=True)] == [ ( (AnyStr("child:"), AnyStr("child_1:")), @@ -9886,6 +11976,7 @@ def test_doubly_nested_graph_state( "parent_2": {"my_key": "hi my value here and there and back again"} }, "step": 3, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -9918,6 +12009,7 @@ def test_doubly_nested_graph_state( "parent_2": {"my_key": "hi my value here and there and back again"} }, "step": 3, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -9943,6 +12035,7 @@ def test_doubly_nested_graph_state( "writes": {"child": {"my_key": "hi my value here and there"}}, "step": 2, "parents": {}, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -9990,6 +12083,7 @@ def test_doubly_nested_graph_state( "source": "loop", "writes": {"parent_1": {"my_key": "hi my value"}}, "step": 1, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -10010,7 +12104,13 @@ def test_doubly_nested_graph_state( "checkpoint_id": AnyStr(), } }, - metadata={"source": "loop", "writes": None, "step": 0, "parents": {}}, + metadata={ + "source": "loop", + "writes": None, + "step": 0, + "parents": {}, + "thread_id": "1", + }, created_at=AnyStr(), parent_config={ "configurable": { @@ -10043,6 +12143,7 @@ def test_doubly_nested_graph_state( "writes": {"__start__": {"my_key": "my value"}}, "step": -1, "parents": {}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=None, @@ -10077,6 +12178,13 @@ def test_doubly_nested_graph_state( "writes": {"child_1": {"my_key": "hi my value here and there"}}, "step": 1, "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child", + "langgraph_path": [PULL, AnyStr("child")], + "langgraph_step": 2, + "langgraph_triggers": [AnyStr("parent_1")], + "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), parent_config={ @@ -10109,6 +12217,13 @@ def test_doubly_nested_graph_state( "writes": None, "step": 0, "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child", + "langgraph_path": [PULL, AnyStr("child")], + "langgraph_step": 2, + "langgraph_triggers": [AnyStr("parent_1")], + "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), parent_config={ @@ -10154,6 +12269,13 @@ def test_doubly_nested_graph_state( "writes": {"__start__": {"my_key": "hi my value"}}, "step": -1, "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child", + "langgraph_path": [PULL, AnyStr("child")], + "langgraph_step": 2, + "langgraph_triggers": [AnyStr("parent_1")], + "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), parent_config=None, @@ -10197,6 +12319,16 @@ def test_doubly_nested_graph_state( AnyStr("child:"): AnyStr(), } ), + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [AnyStr("start:child_1")], }, created_at=AnyStr(), parent_config={ @@ -10242,6 +12374,16 @@ def test_doubly_nested_graph_state( AnyStr("child:"): AnyStr(), } ), + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [AnyStr("start:child_1")], }, created_at=AnyStr(), parent_config={ @@ -10294,6 +12436,16 @@ def test_doubly_nested_graph_state( AnyStr("child:"): AnyStr(), } ), + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [AnyStr("start:child_1")], }, created_at=AnyStr(), parent_config={ @@ -10346,6 +12498,16 @@ def test_doubly_nested_graph_state( AnyStr("child:"): AnyStr(), } ), + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [AnyStr("start:child_1")], }, created_at=AnyStr(), parent_config=None, @@ -10426,13 +12588,37 @@ def test_send_to_nested_graphs( # check state outer_state = graph.get_state(config) + + if not FF_SEND_V2: + # update state of dogs joke graph + graph.update_state(outer_state.tasks[1].state, {"subject": "turtles - hohoho"}) + + # continue past interrupt + assert sorted( + graph.stream(None, config=config), + key=lambda d: d["generate_joke"]["jokes"][0], + ) == [ + {"generate_joke": {"jokes": ["Joke about cats - hohoho"]}}, + {"generate_joke": {"jokes": ["Joke about turtles - hohoho"]}}, + ] + return + assert outer_state == StateSnapshot( values={"subjects": ["cats", "dogs"], "jokes": []}, tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=("__pregel_pull", "__start__"), + error=None, + interrupts=(), + state=None, + result={"subjects": ["cats", "dogs"]}, + ), PregelTask( AnyStr(), "generate_joke", - (PUSH, 0), + (PUSH, ("__pregel_pull", "__start__"), 1, AnyStr()), state={ "configurable": { "thread_id": "1", @@ -10443,7 +12629,7 @@ def test_send_to_nested_graphs( PregelTask( AnyStr(), "generate_joke", - (PUSH, 1), + (PUSH, ("__pregel_pull", "__start__"), 2, AnyStr()), state={ "configurable": { "thread_id": "1", @@ -10460,18 +12646,18 @@ def test_send_to_nested_graphs( "checkpoint_id": AnyStr(), } }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } + metadata={ + "parents": {}, + "source": "input", + "writes": {"__start__": {"subjects": ["cats", "dogs"]}}, + "step": -1, + "thread_id": "1", }, + created_at=AnyStr(), + parent_config=None, ) # check state of each of the inner tasks - assert graph.get_state(outer_state.tasks[0].state) == StateSnapshot( + assert graph.get_state(outer_state.tasks[1].state) == StateSnapshot( values={"subject": "cats - hohoho", "jokes": []}, next=("generate",), config={ @@ -10492,6 +12678,13 @@ def test_send_to_nested_graphs( "source": "loop", "writes": {"edit": None}, "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + "langgraph_checkpoint_ns": AnyStr("generate_joke:"), + "langgraph_node": "generate_joke", + "langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 1, AnyStr()], + "langgraph_step": 0, + "langgraph_triggers": [PUSH], }, created_at=AnyStr(), parent_config={ @@ -10509,7 +12702,7 @@ def test_send_to_nested_graphs( }, tasks=(PregelTask(id=AnyStr(""), name="generate", path=(PULL, "generate")),), ) - assert graph.get_state(outer_state.tasks[1].state) == StateSnapshot( + assert graph.get_state(outer_state.tasks[2].state) == StateSnapshot( values={"subject": "dogs - hohoho", "jokes": []}, next=("generate",), config={ @@ -10530,6 +12723,13 @@ def test_send_to_nested_graphs( "source": "loop", "writes": {"edit": None}, "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + "langgraph_checkpoint_ns": AnyStr("generate_joke:"), + "langgraph_node": "generate_joke", + "langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 2, AnyStr()], + "langgraph_step": 0, + "langgraph_triggers": [PUSH], }, created_at=AnyStr(), parent_config={ @@ -10548,7 +12748,9 @@ def test_send_to_nested_graphs( tasks=(PregelTask(id=AnyStr(""), name="generate", path=(PULL, "generate")),), ) # update state of dogs joke graph - graph.update_state(outer_state.tasks[1].state, {"subject": "turtles - hohoho"}) + graph.update_state( + outer_state.tasks[2 if FF_SEND_V2 else 1].state, {"subject": "turtles - hohoho"} + ) # continue past interrupt assert sorted( @@ -10582,7 +12784,8 @@ def test_send_to_nested_graphs( {"jokes": ["Joke about turtles - hohoho"]}, ] }, - "step": 1, + "step": 0, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -10623,7 +12826,8 @@ def test_send_to_nested_graphs( {"jokes": ["Joke about turtles - hohoho"]}, ] }, - "step": 1, + "step": 0, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -10634,63 +12838,44 @@ def test_send_to_nested_graphs( } }, ), - StateSnapshot( - values={"subjects": ["cats", "dogs"], "jokes": []}, - tasks=( - PregelTask( - AnyStr(), - "generate_joke", - (PUSH, 0), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, - result={"jokes": ["Joke about cats - hohoho"]}, - ), - PregelTask( - AnyStr(), - "generate_joke", - (PUSH, 1), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, - result={"jokes": ["Joke about turtles - hohoho"]}, - ), - ), - next=("generate_joke", "generate_joke"), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), StateSnapshot( values={"jokes": []}, tasks=( PregelTask( - AnyStr(), - "__start__", - (PULL, "__start__"), + id=AnyStr(), + name="__start__", + path=("__pregel_pull", "__start__"), + error=None, + interrupts=(), + state=None, result={"subjects": ["cats", "dogs"]}, ), + PregelTask( + AnyStr(), + "generate_joke", + (PUSH, ("__pregel_pull", "__start__"), 1, AnyStr()), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + } + }, + result={"jokes": ["Joke about cats - hohoho"]}, + ), + PregelTask( + AnyStr(), + "generate_joke", + (PUSH, ("__pregel_pull", "__start__"), 2, AnyStr()), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("generate_joke:"), + } + }, + result={"jokes": ["Joke about turtles - hohoho"]}, + ), ), - next=("__start__",), + next=("__start__", "generate_joke", "generate_joke"), config={ "configurable": { "thread_id": "1", @@ -10703,6 +12888,7 @@ def test_send_to_nested_graphs( "source": "input", "writes": {"__start__": {"subjects": ["cats", "dogs"]}}, "step": -1, + "thread_id": "1", }, created_at=AnyStr(), parent_config=None, @@ -10864,6 +13050,7 @@ def test_weather_subgraph( "writes": {"router_node": {"route": "weather"}}, "step": 1, "parents": {}, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -10951,6 +13138,7 @@ def test_weather_subgraph( "writes": {"router_node": {"route": "weather"}}, "step": 1, "parents": {}, + "thread_id": "14", }, created_at=AnyStr(), parent_config={ @@ -10991,6 +13179,15 @@ def test_weather_subgraph( "writes": {"model_node": {"city": "San Francisco"}}, "step": 1, "parents": {"": AnyStr()}, + "thread_id": "14", + "checkpoint_ns": AnyStr("weather_graph:"), + "langgraph_node": "weather_graph", + "langgraph_path": [PULL, "weather_graph"], + "langgraph_step": 2, + "langgraph_triggers": [ + "branch:router_node:route_after_prediction:weather_graph" + ], + "langgraph_checkpoint_ns": AnyStr("weather_graph:"), }, created_at=AnyStr(), parent_config={ @@ -11041,6 +13238,7 @@ def test_weather_subgraph( "writes": {"router_node": {"route": "weather"}}, "step": 1, "parents": {}, + "thread_id": "14", }, created_at=AnyStr(), parent_config={ @@ -11078,14 +13276,24 @@ def test_weather_subgraph( } }, metadata={ - "source": "update", "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:router_node:route_after_prediction:weather_graph" + ], + "langgraph_checkpoint_ns": AnyStr("weather_graph:"), }, created_at=AnyStr(), parent_config={ @@ -12021,3 +14229,152 @@ def test_debug_nested_subgraphs(): assert stream_task["interrupts"] == history_task.interrupts assert stream_task.get("error") == history_task.error assert stream_task.get("state") == history_task.state + + +def test_add_sequence(): + class State(TypedDict): + foo: Annotated[list[str], operator.add] + bar: str + + def step1(state: State): + return {"foo": ["step1"], "bar": "baz"} + + def step2(state: State): + return {"foo": ["step2"]} + + # test raising if less than 1 steps + with pytest.raises(ValueError): + StateGraph(State).add_sequence([]) + + # test raising if duplicate step names + with pytest.raises(ValueError): + StateGraph(State).add_sequence([step1, step1]) + + with pytest.raises(ValueError): + StateGraph(State).add_sequence([("foo", step1), ("foo", step1)]) + + # test unnamed steps + builder = StateGraph(State) + builder.add_sequence([step1, step2]) + builder.add_edge(START, "step1") + graph = builder.compile() + result = graph.invoke({"foo": []}) + assert result == {"foo": ["step1", "step2"], "bar": "baz"} + stream_chunks = list(graph.stream({"foo": []})) + assert stream_chunks == [ + {"step1": {"foo": ["step1"], "bar": "baz"}}, + {"step2": {"foo": ["step2"]}}, + ] + + # test named steps + builder_named_steps = StateGraph(State) + builder_named_steps.add_sequence([("meow1", step1), ("meow2", step2)]) + builder_named_steps.add_edge(START, "meow1") + graph_named_steps = builder_named_steps.compile() + result = graph_named_steps.invoke({"foo": []}) + stream_chunks = list(graph_named_steps.stream({"foo": []})) + assert result == {"foo": ["step1", "step2"], "bar": "baz"} + assert stream_chunks == [ + {"meow1": {"foo": ["step1"], "bar": "baz"}}, + {"meow2": {"foo": ["step2"]}}, + ] + + builder_named_steps = StateGraph(State) + builder_named_steps.add_sequence( + [ + ("meow1", lambda state: {"foo": ["foo"]}), + ("meow2", lambda state: {"bar": state["foo"][0] + "bar"}), + ], + ) + builder_named_steps.add_edge(START, "meow1") + graph_named_steps = builder_named_steps.compile() + result = graph_named_steps.invoke({"foo": []}) + stream_chunks = list(graph_named_steps.stream({"foo": []})) + # filtered by output schema + assert result == {"bar": "foobar", "foo": ["foo"]} + assert stream_chunks == [ + {"meow1": {"foo": ["foo"]}}, + {"meow2": {"bar": "foobar"}}, + ] + + # test two sequences + + def a(state: State): + return {"foo": ["a"]} + + def b(state: State): + return {"foo": ["b"]} + + builder_two_sequences = StateGraph(State) + builder_two_sequences.add_sequence([a]) + builder_two_sequences.add_sequence([b]) + builder_two_sequences.add_edge(START, "a") + builder_two_sequences.add_edge("a", "b") + graph_two_sequences = builder_two_sequences.compile() + + result = graph_two_sequences.invoke({"foo": []}) + assert result == {"foo": ["a", "b"]} + + stream_chunks = list(graph_two_sequences.stream({"foo": []})) + assert stream_chunks == [ + {"a": {"foo": ["a"]}}, + {"b": {"foo": ["b"]}}, + ] + + # test mixed nodes and sequences + + def c(state: State): + return {"foo": ["c"]} + + def d(state: State): + return {"foo": ["d"]} + + def e(state: State): + return {"foo": ["e"]} + + def foo(state: State): + if state["foo"][0] == "a": + return "d" + else: + return "c" + + builder_complex = StateGraph(State) + builder_complex.add_sequence([a, b]) + builder_complex.add_conditional_edges("b", foo) + builder_complex.add_node(c) + builder_complex.add_sequence([d, e]) + builder_complex.add_edge(START, "a") + graph_complex = builder_complex.compile() + + result = graph_complex.invoke({"foo": []}) + assert result == {"foo": ["a", "b", "d", "e"]} + + result = graph_complex.invoke({"foo": ["start"]}) + assert result == {"foo": ["start", "a", "b", "c"]} + + stream_chunks = list(graph_complex.stream({"foo": []})) + assert stream_chunks == [ + {"a": {"foo": ["a"]}}, + {"b": {"foo": ["b"]}}, + {"d": {"foo": ["d"]}}, + {"e": {"foo": ["e"]}}, + ] + + +def test_runnable_passthrough_node_graph() -> None: + class State(TypedDict): + changeme: str + + async def dummy(state): + return state + + agent = dummy | RunnablePassthrough.assign(prediction=RunnableLambda(lambda x: x)) + + graph_builder = StateGraph(State) + + graph_builder.add_node("agent", agent) + graph_builder.add_edge(START, "agent") + + graph = graph_builder.compile() + + assert graph.get_graph(xray=True).to_json() == graph.get_graph(xray=False).to_json() diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 1c45d414f..a31e444e1 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -1,10 +1,12 @@ import asyncio import operator +import random import re import sys import uuid from collections import Counter from contextlib import asynccontextmanager, contextmanager +from dataclasses import replace from time import perf_counter from typing import ( Annotated, @@ -50,11 +52,17 @@ from langgraph.checkpoint.base import ( CheckpointTuple, ) from langgraph.checkpoint.memory import MemorySaver -from langgraph.constants import ERROR, PULL, PUSH +from langgraph.constants import ( + CONFIG_KEY_NODE_FINISHED, + ERROR, + FF_SEND_V2, + PULL, + PUSH, + START, +) from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt -from langgraph.graph import END, Graph, StateGraph -from langgraph.graph.graph import START -from langgraph.graph.message import MessageGraph, add_messages +from langgraph.graph import END, Graph, GraphCommand, StateGraph +from langgraph.graph.message import MessageGraph, MessagesState, add_messages from langgraph.managed.shared_value import SharedValue from langgraph.prebuilt.chat_agent_executor import create_tool_calling_executor from langgraph.prebuilt.tool_node import ToolNode @@ -62,7 +70,14 @@ from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.retry import RetryPolicy from langgraph.store.base import BaseStore from langgraph.store.memory import InMemoryStore -from langgraph.types import Interrupt, PregelTask, Send, StreamWriter +from langgraph.types import ( + Command, + Interrupt, + PregelTask, + Send, + StreamWriter, + interrupt, +) from tests.any_str import AnyDict, AnyStr, AnyVersion, FloatBetween, UnsortedSequence from tests.conftest import ( ALL_CHECKPOINTERS_ASYNC, @@ -243,6 +258,10 @@ async def test_node_cancellation_on_other_node_exception_two() -> None: await graph.ainvoke(1) +@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_dynamic_interrupt(checkpointer_name: str) -> None: class State(TypedDict): @@ -255,8 +274,10 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None: nonlocal tool_two_node_count tool_two_node_count += 1 if s["market"] == "DE": - raise NodeInterrupt("Just because...") - return {"my_key": " all good"} + answer = interrupt("Just because...") + else: + answer = " all good" + return {"my_key": answer} tool_two_graph = StateGraph(State) tool_two_graph.add_node("tool_two", tool_two_node, retry=RetryPolicy()) @@ -289,6 +310,33 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None: 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 [ @@ -297,7 +345,15 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None: {"my_key": "value ⛰️", "market": "DE"}, thread1 ) ] == [ - {"__interrupt__": [Interrupt(value="Just because...", when="during")]}, + { + "__interrupt__": ( + Interrupt( + value="Just because...", + resumable=True, + ns=[AnyStr("tool_two:")], + ), + ) + }, ] assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ { @@ -305,12 +361,14 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None: "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) @@ -322,24 +380,237 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None: AnyStr(), "tool_two", (PULL, "tool_two"), - interrupts=(Interrupt("Just because..."),), + 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}, + 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, + ) + + # 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, ) +@pytest.mark.skipif(not FF_SEND_V2, reason="send v2 is not enabled") +@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_copy_checkpoint(checkpointer_name: str) -> None: + class State(TypedDict): + my_key: Annotated[str, operator.add] + market: str + + def tool_one(s: State) -> State: + return {"my_key": " one"} + + tool_two_node_count = 0 + + def tool_two_node(s: State) -> State: + nonlocal tool_two_node_count + tool_two_node_count += 1 + if s["market"] == "DE": + answer = interrupt("Just because...") + else: + answer = " all good" + return {"my_key": answer} + + def start(state: State) -> list[Union[Send, str]]: + return ["tool_two", Send("tool_one", state)] + + tool_two_graph = StateGraph(State) + tool_two_graph.add_node("tool_two", tool_two_node, retry=RetryPolicy()) + tool_two_graph.add_node("tool_one", tool_one) + tool_two_graph.set_conditional_entry_point(start) + tool_two = tool_two_graph.compile() + + tracer = FakeTracer() + assert await tool_two.ainvoke( + {"my_key": "value", "market": "DE"}, {"callbacks": [tracer]} + ) == { + "my_key": "value one", + "market": "DE", + } + assert tool_two_node_count == 1, "interrupts aren't retried" + assert len(tracer.runs) == 1 + run = tracer.runs[0] + assert run.end_time is not None + assert run.error is None + assert run.outputs == {"market": "DE", "my_key": "value one"} + + assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == { + "my_key": "value one all good", + "market": "US", + } + + async with awith_checkpointer(checkpointer_name) as checkpointer: + tool_two = tool_two_graph.compile(checkpointer=checkpointer) + + # 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 + ) + ] == [ + { + "tool_one": {"my_key": " one"}, + }, + { + "__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 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", + } + assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ + { + "parents": {}, + "source": "loop", + "step": 0, + "writes": {"tool_one": {"my_key": " one"}}, + "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(), + "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": {"tool_one": {"my_key": " one"}}, + "thread_id": "1", + }, + parent_config=[ + c async for c in tool_two.checkpointer.alist(thread1, limit=2) + ][-1].config, + ) + # clear the interrupt and next tasks + await tool_two.aupdate_state(thread1, None) + # 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 ⛰️ one", "market": "DE"}, + next=("tool_two",), + tasks=( + PregelTask( + AnyStr(), + "tool_two", + (PULL, "tool_two"), + interrupts=(), + ), + ), + 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, + ) + + +@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_node_not_cancelled_on_other_node_interrupted( checkpointer_name: str, ) -> None: class State(TypedDict): - hello: str + hello: Annotated[str, operator.add] awhiles = 0 inner_task_cancelled = False @@ -350,15 +621,14 @@ async def test_node_not_cancelled_on_other_node_interrupted( awhiles += 1 try: await asyncio.sleep(1) - return {"hello": "again"} + return {"hello": " again"} except asyncio.CancelledError: nonlocal inner_task_cancelled inner_task_cancelled = True raise async def iambad(input: State) -> None: - if input["hello"] != "bye": - raise NodeInterrupt("I am bad") + return {"hello": interrupt("I am bad")} builder = StateGraph(State) builder.add_node("agent", awhile) @@ -369,20 +639,26 @@ async def test_node_not_cancelled_on_other_node_interrupted( graph = builder.compile(checkpointer=checkpointer) thread = {"configurable": {"thread_id": "1"}} - assert await graph.ainvoke({"hello": "world"}, thread) == {"hello": "world"} + # writes from "awhile" are applied to last chunk + assert await graph.ainvoke({"hello": "world"}, thread) == { + "hello": "world again" + } assert not inner_task_cancelled assert awhiles == 1 - assert await graph.ainvoke(None, thread, debug=True) == {"hello": "world"} + assert await graph.ainvoke(None, thread, debug=True) == {"hello": "world again"} assert not inner_task_cancelled assert awhiles == 1 - assert await graph.ainvoke({"hello": "bye"}, thread) == {"hello": "again"} + # resume with answer + assert await graph.ainvoke(Command(resume=" okay"), thread) == { + "hello": "world again okay" + } assert not inner_task_cancelled - assert awhiles == 2 + assert awhiles == 1 @pytest.mark.repeat(10) @@ -473,19 +749,18 @@ async def test_cancel_graph_astream(checkpointer_name: str) -> None: assert awhile.started is False # checkpoint with output of "alittlewhile" should not be saved + # but we should have applied pending writes if checkpointer is not None: state = await graph.aget_state(thread1) assert state is not None - assert state.values == {"value": 1} - assert state.next == ( - "aparallelwhile", - "alittlewhile", - ) + assert state.values == {"value": 3} # 1 + 2 + assert state.next == ("aparallelwhile",) assert state.metadata == { "parents": {}, "source": "loop", "step": 0, "writes": None, + "thread_id": "1", } @@ -562,6 +837,7 @@ async def test_cancel_graph_astream_events_v2(checkpointer_name: Optional[str]) "source": "loop", "step": 1, "writes": {"alittlewhile": {"value": 2}}, + "thread_id": "2", } @@ -945,6 +1221,7 @@ async def test_invoke_two_processes_in_out_interrupt( "source": "loop", "step": 6, "writes": {"two": 5}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=history[1].config, @@ -967,6 +1244,7 @@ async def test_invoke_two_processes_in_out_interrupt( "source": "loop", "step": 5, "writes": {"one": None}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=history[2].config, @@ -989,6 +1267,7 @@ async def test_invoke_two_processes_in_out_interrupt( "source": "input", "step": 4, "writes": {"input": 3}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=history[3].config, @@ -1009,6 +1288,7 @@ async def test_invoke_two_processes_in_out_interrupt( "source": "loop", "step": 3, "writes": {"one": None}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=history[4].config, @@ -1031,6 +1311,7 @@ async def test_invoke_two_processes_in_out_interrupt( "source": "input", "step": 2, "writes": {"input": 20}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=history[5].config, @@ -1051,6 +1332,7 @@ async def test_invoke_two_processes_in_out_interrupt( "source": "loop", "step": 1, "writes": {"two": 4}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=history[6].config, @@ -1073,6 +1355,7 @@ async def test_invoke_two_processes_in_out_interrupt( "source": "loop", "step": 0, "writes": {"one": None}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=history[7].config, @@ -1095,6 +1378,7 @@ async def test_invoke_two_processes_in_out_interrupt( "source": "input", "step": -1, "writes": {"input": 2}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=None, @@ -1170,6 +1454,7 @@ async def test_fork_always_re_runs_nodes( "source": "loop", "step": 5, "writes": {"add_one": 1}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=history[1].config, @@ -1190,6 +1475,7 @@ async def test_fork_always_re_runs_nodes( "source": "loop", "step": 4, "writes": {"add_one": 1}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=history[2].config, @@ -1210,6 +1496,7 @@ async def test_fork_always_re_runs_nodes( "source": "loop", "step": 3, "writes": {"add_one": 1}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=history[3].config, @@ -1230,6 +1517,7 @@ async def test_fork_always_re_runs_nodes( "source": "loop", "step": 2, "writes": {"add_one": 1}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=history[4].config, @@ -1250,6 +1538,7 @@ async def test_fork_always_re_runs_nodes( "source": "loop", "step": 1, "writes": {"add_one": 1}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=history[5].config, @@ -1265,7 +1554,13 @@ async def test_fork_always_re_runs_nodes( "checkpoint_id": AnyStr(), } }, - metadata={"parents": {}, "source": "loop", "step": 0, "writes": None}, + metadata={ + "parents": {}, + "source": "loop", + "step": 0, + "writes": None, + "thread_id": "1", + }, created_at=AnyStr(), parent_config=history[6].config, ), @@ -1287,6 +1582,7 @@ async def test_fork_always_re_runs_nodes( "source": "input", "step": -1, "writes": {"__start__": 1}, + "thread_id": "1", }, created_at=AnyStr(), parent_config=None, @@ -1691,10 +1987,11 @@ async def test_pending_writes_resume( 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": 1} - assert state.next == ("one", "two") + assert state.values == {"value": 3} + assert state.next == ("two",) assert state.tasks == ( PregelTask(AnyStr(), "one", (PULL, "one"), result={"value": 2}), PregelTask( @@ -1709,7 +2006,13 @@ async def test_pending_writes_resume( "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 checkpointer.aget_tuple(thread1) assert checkpoint is not None @@ -1797,6 +2100,7 @@ async def test_pending_writes_resume( "step": 1, "source": "loop", "writes": {"one": {"value": 2}, "two": {"value": 3}}, + "thread_id": "1", }, parent_config={ "configurable": { @@ -1843,7 +2147,13 @@ async def test_pending_writes_resume( "start:two": "__start__", }, }, - metadata={"parents": {}, "step": 0, "source": "loop", "writes": None}, + metadata={ + "parents": {}, + "step": 0, + "source": "loop", + "writes": None, + "thread_id": "1", + }, parent_config={ "configurable": { "thread_id": "1", @@ -1885,6 +2195,7 @@ async def test_pending_writes_resume( "step": -1, "source": "input", "writes": {"__start__": {"value": 1}}, + "thread_id": "1", }, parent_config=None, pending_writes=UnsortedSequence( @@ -1895,6 +2206,85 @@ async def test_pending_writes_resume( ) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_run_from_checkpoint_id_retains_previous_writes( + request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture +) -> None: + class MyState(TypedDict): + myval: Annotated[int, operator.add] + otherval: bool + + class Anode: + def __init__(self): + self.switch = False + + async def __call__(self, state: MyState): + self.switch = not self.switch + return {"myval": 2 if self.switch else 1, "otherval": self.switch} + + builder = StateGraph(MyState) + thenode = Anode() # Fun. + builder.add_node("node_one", thenode) + builder.add_node("node_two", thenode) + builder.add_edge(START, "node_one") + + def _getedge(src: str): + swap = "node_one" if src == "node_two" else "node_two" + + def _edge(st: MyState) -> Literal["__end__", "node_one", "node_two"]: + if st["myval"] > 3: + return END + if st["otherval"]: + return swap + return src + + return _edge + + builder.add_conditional_edges("node_one", _getedge("node_one")) + builder.add_conditional_edges("node_two", _getedge("node_two")) + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer) + + thread_id = uuid.uuid4() + thread1 = {"configurable": {"thread_id": str(thread_id)}} + + result = await graph.ainvoke({"myval": 1}, thread1) + assert result["myval"] == 4 + history = [c async for c in graph.aget_state_history(thread1)] + + assert len(history) == 4 + assert history[-1].values == {"myval": 0} + assert history[0].values == {"myval": 4, "otherval": False} + + second_run_config = { + **thread1, + "configurable": { + **thread1["configurable"], + "checkpoint_id": history[1].config["configurable"]["checkpoint_id"], + }, + } + second_result = await graph.ainvoke(None, second_run_config) + assert second_result == {"myval": 5, "otherval": True} + + new_history = [ + c + async for c in graph.aget_state_history( + {"configurable": {"thread_id": str(thread_id), "checkpoint_ns": ""}} + ) + ] + + assert len(new_history) == len(history) + 1 + for original, new in zip(history, new_history[1:]): + assert original.values == new.values + assert original.next == new.next + assert original.metadata["step"] == new.metadata["step"] + + def _get_tasks(hist: list, start: int): + return [h.tasks for h in hist[start:]] + + assert _get_tasks(new_history, 1) == _get_tasks(history, 0) + + async def test_cond_edge_after_send() -> None: class Node: def __init__(self, name: str): @@ -1922,6 +2312,1334 @@ async def test_cond_edge_after_send() -> None: assert await graph.ainvoke(["0"]) == ["0", "1", "2", "2", "3"] +async def test_concurrent_emit_sends() -> None: + class Node: + def __init__(self, name: str): + self.name = name + setattr(self, "__name__", name) + + async def __call__(self, state): + return ( + [self.name] + if isinstance(state, list) + else ["|".join((self.name, str(state)))] + ) + + async def send_for_fun(state): + return [Send("2", 1), Send("2", 2), "3.1"] + + async def send_for_profit(state): + return [Send("2", 3), Send("2", 4)] + + async def route_to_three(state) -> Literal["3"]: + return "3" + + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node(Node("1")) + builder.add_node(Node("1.1")) + builder.add_node(Node("2")) + builder.add_node(Node("3")) + builder.add_node(Node("3.1")) + builder.add_edge(START, "1") + builder.add_edge(START, "1.1") + builder.add_conditional_edges("1", send_for_fun) + builder.add_conditional_edges("1.1", send_for_profit) + builder.add_conditional_edges("2", route_to_three) + graph = builder.compile() + assert await graph.ainvoke(["0"]) == ( + [ + "0", + "1", + "1.1", + "2|1", + "2|2", + "2|3", + "2|4", + "3", + "3.1", + ] + if FF_SEND_V2 + else [ + "0", + "1", + "1.1", + "3.1", + "2|1", + "2|2", + "2|3", + "2|4", + "3", + ] + ) + + +@pytest.mark.repeat(10) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_send_sequences(checkpointer_name: str) -> None: + class Node: + def __init__(self, name: str): + self.name = name + setattr(self, "__name__", name) + + async def __call__(self, state): + update = ( + [self.name] + if isinstance(state, list) # or isinstance(state, Control) + else ["|".join((self.name, str(state)))] + ) + if isinstance(state, GraphCommand): + return replace(state, update=update) + else: + return update + + async def send_for_fun(state): + return [ + Send("2", GraphCommand(send=Send("2", 3))), + Send("2", GraphCommand(send=Send("2", 4))), + "3.1", + ] + + async def route_to_three(state) -> Literal["3"]: + return "3" + + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node(Node("1")) + builder.add_node(Node("2")) + builder.add_node(Node("3")) + builder.add_node(Node("3.1")) + builder.add_edge(START, "1") + builder.add_conditional_edges("1", send_for_fun) + builder.add_conditional_edges("2", route_to_three) + graph = builder.compile() + assert ( + await graph.ainvoke(["0"]) + == [ + "0", + "1", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='2', arg=4))", + "2|3", + "2|4", + "3", + "3.1", + ] + if FF_SEND_V2 + else [ + "0", + "1", + "3.1", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='2', arg=4))", + "3", + "2|3", + "2|4", + "3", + ] + ) + + if not FF_SEND_V2: + return + + 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", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='2', arg=4))", + "2|3", + "2|4", + ] + assert await graph.ainvoke(None, thread1) == [ + "0", + "1", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='2', arg=4))", + "2|3", + "2|4", + "3", + "3.1", + ] + + +@pytest.mark.repeat(20) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: + if not FF_SEND_V2: + pytest.skip("Send deduplication is only available in Send V2") + + class InterruptOnce: + ticks: int = 0 + + def __call__(self, state): + self.ticks += 1 + if self.ticks == 1: + raise NodeInterrupt("Bahh") + return ["|".join(("flaky", str(state)))] + + class Node: + def __init__(self, name: str): + self.name = name + self.ticks = 0 + setattr(self, "__name__", name) + + def __call__(self, state): + self.ticks += 1 + update = ( + [self.name] + if isinstance(state, list) + else ["|".join((self.name, str(state)))] + ) + if isinstance(state, GraphCommand): + return replace(state, update=update) + else: + return update + + def send_for_fun(state): + return [ + Send("2", GraphCommand(send=Send("2", 3))), + Send("2", GraphCommand(send=Send("flaky", 4))), + "3.1", + ] + + def route_to_three(state) -> Literal["3"]: + return "3" + + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node(Node("1")) + builder.add_node(Node("2")) + builder.add_node(Node("3")) + builder.add_node(Node("3.1")) + builder.add_node("flaky", InterruptOnce()) + builder.add_edge(START, "1") + builder.add_conditional_edges("1", send_for_fun) + builder.add_conditional_edges("2", route_to_three) + + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer) + thread1 = {"configurable": {"thread_id": "1"}} + assert await graph.ainvoke(["0"], thread1, debug=1) == [ + "0", + "1", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='flaky', arg=4))", + "2|3", + ] + assert builder.nodes["2"].runnable.func.ticks == 3 + assert builder.nodes["flaky"].runnable.func.ticks == 1 + # resume execution + assert await graph.ainvoke(None, thread1, debug=1) == [ + "0", + "1", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='flaky', arg=4))", + "2|3", + "flaky|4", + "3", + "3.1", + ] + # node "2" doesn't get called again, as we recover writes saved before + assert builder.nodes["2"].runnable.func.ticks == 3 + # node "flaky" gets called again, as it was interrupted + assert builder.nodes["flaky"].runnable.func.ticks == 2 + # check history + history = [c async for c in graph.aget_state_history(thread1)] + assert history == [ + StateSnapshot( + values=[ + "0", + "1", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='flaky', arg=4))", + "2|3", + "flaky|4", + "3", + "3.1", + ], + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"3": ["3"], "3.1": ["3.1"]}, + "thread_id": "1", + "step": 2, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=(), + ), + StateSnapshot( + values=[ + "0", + "1", + "2|Command(send=Send(node='2', arg=3))", + "2|Command(send=Send(node='flaky', arg=4))", + "2|3", + "flaky|4", + ], + next=("3", "3.1"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "1": ["1"], + "2": [ + ["2|Command(send=Send(node='2', arg=3))"], + ["2|Command(send=Send(node='flaky', arg=4))"], + ["2|3"], + ], + "flaky": ["flaky|4"], + }, + "thread_id": "1", + "step": 1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="3", + path=("__pregel_pull", "3"), + error=None, + interrupts=(), + state=None, + result=["3"], + ), + PregelTask( + id=AnyStr(), + name="3.1", + path=("__pregel_pull", "3.1"), + error=None, + interrupts=(), + state=None, + result=["3.1"], + ), + ), + ), + StateSnapshot( + values=["0"], + next=("1", "2", "2", "2", "flaky"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": None, + "thread_id": "1", + "step": 0, + "parents": {}, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="1", + path=("__pregel_pull", "1"), + error=None, + interrupts=(), + state=None, + result=["1"], + ), + PregelTask( + id=AnyStr(), + name="2", + path=( + "__pregel_push", + ("__pregel_pull", "1"), + 2, + AnyStr(), + ), + error=None, + interrupts=(), + state=None, + result=["2|Command(send=Send(node='2', arg=3))"], + ), + PregelTask( + id=AnyStr(), + name="2", + path=( + "__pregel_push", + ("__pregel_pull", "1"), + 3, + AnyStr(), + ), + error=None, + interrupts=(), + state=None, + result=["2|Command(send=Send(node='flaky', arg=4))"], + ), + PregelTask( + id=AnyStr(), + name="2", + path=( + "__pregel_push", + ( + "__pregel_push", + ("__pregel_pull", "1"), + 2, + AnyStr(), + ), + 2, + AnyStr(), + ), + error=None, + interrupts=(), + state=None, + result=["2|3"], + ), + PregelTask( + id=AnyStr(), + name="flaky", + path=( + "__pregel_push", + ( + "__pregel_push", + ("__pregel_pull", "1"), + 3, + AnyStr(), + ), + 2, + AnyStr(), + ), + error=None, + interrupts=(Interrupt(value="Bahh", when="during"),), + state=None, + result=["flaky|4"], + ), + ), + ), + StateSnapshot( + values=[], + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"__start__": ["0"]}, + "thread_id": "1", + "step": -1, + "parents": {}, + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=("__pregel_pull", "__start__"), + error=None, + interrupts=(), + state=None, + result=["0"], + ), + ), + ), + ] + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_send_react_interrupt(checkpointer_name: str) -> None: + from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage + + ai_message = AIMessage( + "", + id="ai1", + tool_calls=[ToolCall(name="foo", args={"hi": [1, 2, 3]}, id=AnyStr())], + ) + + async def agent(state): + return {"messages": ai_message} + + def route(state): + if isinstance(state["messages"][-1], AIMessage): + return [ + Send(call["name"], call) for call in state["messages"][-1].tool_calls + ] + + foo_called = 0 + + async def foo(call: ToolCall): + nonlocal foo_called + foo_called += 1 + return {"messages": ToolMessage(str(call["args"]), tool_call_id=call["id"])} + + builder = StateGraph(MessagesState) + builder.add_node(agent) + builder.add_node(foo) + builder.add_edge(START, "agent") + builder.add_conditional_edges("agent", route) + graph = builder.compile() + + assert await graph.ainvoke({"messages": [HumanMessage("hello")]}) == { + "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 + + 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 + + # 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) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + } + assert foo_called == 0 + + if not FF_SEND_V2: + return + + # get state should show the pending task + state = await graph.aget_state(thread1) + assert state == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + }, + next=("foo",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 0, + "source": "loop", + "writes": None, + "parents": {}, + "thread_id": "2", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + content="", + additional_kwargs={}, + response_metadata={}, + id="ai1", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ) + }, + ), + PregelTask( + id=AnyStr(), + name="foo", + path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()), + error=None, + interrupts=(), + state=None, + result=None, + ), + ), + ) + + # 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": 1, + "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=(), + ) + + # 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) == { + "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( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + }, + next=("foo",), + config={ + "configurable": { + "thread_id": "3", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 0, + "source": "loop", + "writes": None, + "parents": {}, + "thread_id": "3", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "3", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + "", + id="ai1", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ) + }, + ), + PregelTask( + id=AnyStr(), + name="foo", + path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()), + error=None, + interrupts=(), + state=None, + result=None, + ), + ), + ) + + # 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", + "args": {"hi": [4, 5, 6]}, + "id": "tool1", + "type": "tool_call", + } + ], + ), + ] + }, + next=("foo",), + config={ + "configurable": { + "thread_id": "3", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 1, + "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, AnyStr()), + error=None, + interrupts=(), + state=None, + result=None, + ), + ), + ) + + # 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 + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_send_react_interrupt_control( + checkpointer_name: str, snapshot: SnapshotAssertion +) -> None: + from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage + + ai_message = AIMessage( + "", + id="ai1", + tool_calls=[ToolCall(name="foo", args={"hi": [1, 2, 3]}, id=AnyStr())], + ) + + async def agent(state) -> Command[Literal["foo"]]: + return GraphCommand( + update={"messages": ai_message}, + send=[Send(call["name"], call) for call in ai_message.tool_calls], + ) + + foo_called = 0 + + async def foo(call: ToolCall): + nonlocal foo_called + foo_called += 1 + return {"messages": ToolMessage(str(call["args"]), tool_call_id=call["id"])} + + builder = StateGraph(MessagesState) + builder.add_node(agent) + builder.add_node(foo) + builder.add_edge(START, "agent") + graph = builder.compile() + assert graph.get_graph().draw_mermaid() == snapshot + + assert await graph.ainvoke({"messages": [HumanMessage("hello")]}) == { + "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 + + 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 + + # 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) == { + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + } + assert foo_called == 0 + + if not FF_SEND_V2: + return + + # get state should show the pending task + state = await graph.aget_state(thread1) + assert state == StateSnapshot( + values={ + "messages": [ + _AnyIdHumanMessage(content="hello"), + _AnyIdAIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ), + ] + }, + next=("foo",), + config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "step": 0, + "source": "loop", + "writes": None, + "parents": {}, + "thread_id": "2", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "2", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + content="", + additional_kwargs={}, + response_metadata={}, + id="ai1", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ) + }, + ), + PregelTask( + id=AnyStr(), + name="foo", + path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()), + error=None, + interrupts=(), + state=None, + result=None, + ), + ), + ) + + # 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": 1, + "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=(), + ) + + # 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: + class Node: + def __init__(self, name: str): + self.name = name + setattr(self, "__name__", name) + self.currently = 0 + self.max_currently = 0 + + async def __call__(self, state): + self.currently += 1 + if self.currently > self.max_currently: + self.max_currently = self.currently + await asyncio.sleep(random.random() / 10) + self.currently -= 1 + return [state] + + def one(state): + return ["1"] + + def three(state): + return ["3"] + + async def send_to_many(state): + return [Send("2", idx) for idx in range(100)] + + async def route_to_three(state) -> Literal["3"]: + return "3" + + node2 = Node("2") + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node("1", one) + builder.add_node(node2) + builder.add_node("3", three) + builder.add_edge(START, "1") + builder.add_conditional_edges("1", send_to_many) + builder.add_conditional_edges("2", route_to_three) + graph = builder.compile() + + assert await graph.ainvoke(["0"]) == ["0", "1", *range(100), "3"] + assert node2.max_currently == 100 + assert node2.currently == 0 + node2.max_currently = 0 + + assert await graph.ainvoke(["0"], {"max_concurrency": 10}) == [ + "0", + "1", + *range(100), + "3", + ] + 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"}} + + 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 node1(state) -> GraphCommand[Literal["2"]]: + return GraphCommand(update=["1"], send=[Send("2", idx) for idx in range(100)]) + + node2_currently = 0 + node2_max_currently = 0 + + async def node2(state) -> GraphCommand[Literal["3"]]: + nonlocal node2_currently, node2_max_currently + node2_currently += 1 + if node2_currently > node2_max_currently: + node2_max_currently = node2_currently + await asyncio.sleep(0.1) + node2_currently -= 1 + + return GraphCommand(update=[state], goto="3") + + async def node3(state) -> Literal["3"]: + return ["3"] + + builder = StateGraph(Annotated[list, operator.add]) + builder.add_node("1", node1) + builder.add_node("2", node2) + builder.add_node("3", node3) + builder.add_edge(START, "1") + graph = builder.compile() + + assert ( + graph.get_graph().draw_mermaid() + == """%%{init: {'flowchart': {'curve': 'linear'}}}%% +graph TD; + __start__([__start__
]):::first + 1(1) + 2(2) + 3([3]):::last + __start__ --> 1; + 1 -.-> 2; + 2 -.-> 3; + classDef default fill:#f2f0ff,line-height:1.2 + classDef first fill-opacity:0 + classDef last fill:#bfb6fc +""" + ) + + assert await graph.ainvoke(["0"], debug=True) == ["0", "1", *range(100), "3"] + assert node2_max_currently == 100 + assert node2_currently == 0 + node2_max_currently = 0 + + assert await graph.ainvoke(["0"], {"max_concurrency": 10}) == [ + "0", + "1", + *range(100), + "3", + ] + 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"}} + + 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 @@ -2603,6 +4321,7 @@ async def test_conditional_graph(checkpointer_name: str) -> None: } } }, + "thread_id": "1", }, parent_config=[ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) @@ -2652,6 +4371,7 @@ async def test_conditional_graph(checkpointer_name: str) -> None: "input": "what is weather in sf", } }, + "thread_id": "1", }, parent_config=[ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) @@ -2776,6 +4496,7 @@ async def test_conditional_graph(checkpointer_name: str) -> None: ), } }, + "thread_id": "1", }, parent_config=[ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) @@ -2842,6 +4563,7 @@ async def test_conditional_graph(checkpointer_name: str) -> None: } } }, + "thread_id": "2", }, parent_config=[ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) @@ -2891,6 +4613,7 @@ async def test_conditional_graph(checkpointer_name: str) -> None: "input": "what is weather in sf", } }, + "thread_id": "2", }, parent_config=[ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) @@ -3015,6 +4738,7 @@ async def test_conditional_graph(checkpointer_name: str) -> None: ), } }, + "thread_id": "2", }, parent_config=[ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) @@ -3081,6 +4805,7 @@ async def test_conditional_graph(checkpointer_name: str) -> None: } } }, + "thread_id": "3", }, parent_config=[ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) @@ -3501,6 +5226,7 @@ async def test_conditional_graph_state( ), } }, + "thread_id": "1", }, parent_config=[ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) @@ -3547,6 +5273,7 @@ async def test_conditional_graph_state( ) } }, + "thread_id": "1", }, parent_config=[ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) @@ -3627,6 +5354,7 @@ async def test_conditional_graph_state( ) } }, + "thread_id": "1", }, parent_config=[ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) @@ -3686,6 +5414,7 @@ async def test_conditional_graph_state( ), } }, + "thread_id": "2", }, parent_config=[ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) @@ -3731,6 +5460,7 @@ async def test_conditional_graph_state( ) } }, + "thread_id": "2", }, parent_config=[ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) @@ -3809,6 +5539,7 @@ async def test_conditional_graph_state( ) } }, + "thread_id": "2", }, parent_config=[ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) @@ -3999,18 +5730,6 @@ async def test_prebuilt_tool_chat() -> None: stream_mode="messages", ) ] == [ - ( - _AnyIdHumanMessage( - content="what is weather in sf", - ), - { - "langgraph_step": 0, - "langgraph_node": "__start__", - "langgraph_triggers": ["__start__"], - "langgraph_path": ("__pregel_pull", "__start__"), - "langgraph_checkpoint_ns": AnyStr("__start__:"), - }, - ), ( _AnyIdAIMessageChunk( content="", @@ -4443,6 +6162,8 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ] async with awith_checkpointer(checkpointer_name) as checkpointer: + # interrupt after agent + app_w_interrupt = workflow.compile( checkpointer=checkpointer, interrupt_after=["agent"], @@ -4473,6 +6194,9 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: {"__interrupt__": ()}, ] + if not FF_SEND_V2: + return + assert await app_w_interrupt.aget_state(config) == StateSnapshot( values={ "messages": [ @@ -4490,7 +6214,33 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ), ] }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),), + tasks=( + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + "", + id="ai1", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "query"}, + "id": "tool_call123", + "type": "tool_call", + } + ], + ) + }, + ), + PregelTask( + AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) + ), + ), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=( @@ -4499,22 +6249,9 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: metadata={ "parents": {}, "source": "loop", - "step": 1, - "writes": { - "agent": { - "messages": AIMessage( - id="ai1", - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ) - } - }, + "step": 0, + "writes": None, + "thread_id": "1", }, parent_config=[ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) @@ -4545,14 +6282,14 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ), ] }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),), + tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0, AnyStr())),), next=("tools",), config=tup.config, created_at=tup.checkpoint["ts"], metadata={ "parents": {}, "source": "update", - "step": 2, + "step": 1, "writes": { "agent": { "messages": AIMessage( @@ -4568,6 +6305,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ) } }, + "thread_id": "1", }, parent_config=[ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) @@ -4647,8 +6385,40 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ] }, tasks=( - PregelTask(AnyStr(), "tools", (PUSH, 0)), - PregelTask(AnyStr(), "tools", (PUSH, 1)), + PregelTask( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + "", + id="ai2", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "another", "idx": 0}, + "id": "tool_call234", + "type": "tool_call", + }, + { + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + "id": "tool_call567", + "type": "tool_call", + }, + ], + ) + }, + ), + PregelTask( + AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) + ), + PregelTask( + AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3, AnyStr()) + ), ), next=("tools", "tools"), config=tup.config, @@ -4656,27 +6426,17 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: metadata={ "parents": {}, "source": "loop", - "step": 4, + "step": 2, "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}, - }, - ], - ) + "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) @@ -4720,12 +6480,346 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: metadata={ "parents": {}, "source": "update", - "step": 5, + "step": 3, "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, + ) + + # 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 + ) + ] == [ + { + "agent": { + "messages": AIMessage( + id="ai1", + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ) + } + }, + {"__interrupt__": ()}, + ] + + 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( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + content="", + additional_kwargs={}, + response_metadata={}, + id="ai1", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "query"}, + "id": "tool_call123", + "type": "tool_call", + } + ], + ) + }, + ), + PregelTask( + AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) + ), + ), + 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": None, + "thread_id": "2", + }, + 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["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=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "a different query"}, + }, + ], + ), + ] + }, + tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0, AnyStr())),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "update", + "step": 1, + "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, + ) + + 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( + id=AnyStr(), + name="agent", + path=("__pregel_pull", "agent"), + error=None, + interrupts=(), + state=None, + result={ + "messages": AIMessage( + content="", + additional_kwargs={}, + response_metadata={}, + id="ai2", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "another", "idx": 0}, + "id": "tool_call234", + "type": "tool_call", + }, + { + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + "id": "tool_call567", + "type": "tool_call", + }, + ], + ) + }, + ), + PregelTask( + AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) + ), + PregelTask( + AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3, AnyStr()) + ), + ), + next=("tools", "tools"), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 2, + "writes": { + "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, + ) + + 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": 3, + "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) @@ -4986,6 +7080,7 @@ async def test_message_graph(checkpointer_name: str) -> None: id="ai1", ) }, + "thread_id": "1", }, parent_config=[ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) @@ -5035,6 +7130,7 @@ async def test_message_graph(checkpointer_name: str) -> None: id="ai1", ) }, + "thread_id": "1", }, parent_config=[ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) @@ -5120,6 +7216,7 @@ async def test_message_graph(checkpointer_name: str) -> None: id="ai2", ) }, + "thread_id": "1", }, parent_config=[ c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) @@ -5163,6 +7260,7 @@ async def test_message_graph(checkpointer_name: str) -> None: "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) @@ -5465,12 +7563,16 @@ async def test_start_branch_then(checkpointer_name: str) -> None: "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( @@ -5481,7 +7583,14 @@ async def test_start_branch_then(checkpointer_name: str) -> None: created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ "ts" ], - metadata={"parents": {}, "source": "loop", "step": 0, "writes": None}, + 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, @@ -5504,6 +7613,8 @@ async def test_start_branch_then(checkpointer_name: str) -> None: "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) @@ -5524,7 +7635,14 @@ async def test_start_branch_then(checkpointer_name: str) -> None: created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ "ts" ], - metadata={"parents": {}, "source": "loop", "step": 0, "writes": None}, + 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, @@ -5547,6 +7665,8 @@ async def test_start_branch_then(checkpointer_name: str) -> None: "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) @@ -5567,7 +7687,14 @@ async def test_start_branch_then(checkpointer_name: str) -> None: created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ "ts" ], - metadata={"parents": {}, "source": "loop", "step": 0, "writes": None}, + 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, @@ -5587,6 +7714,8 @@ async def test_start_branch_then(checkpointer_name: str) -> None: "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) @@ -5610,6 +7739,8 @@ async def test_start_branch_then(checkpointer_name: str) -> None: "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) @@ -5678,6 +7809,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "source": "input", "step": -1, "writes": {"__start__": {"my_key": "value", "market": "DE"}}, + "thread_id": "10", }, "parent_config": None, "next": ["__start__"], @@ -5716,6 +7848,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "source": "loop", "step": 0, "writes": None, + "thread_id": "10", }, "parent_config": { "tags": [], @@ -5787,6 +7920,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "source": "loop", "step": 1, "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "10", }, "parent_config": { "tags": [], @@ -5858,6 +7992,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "source": "loop", "step": 2, "writes": {"tool_two_slow": {"my_key": " slow"}}, + "thread_id": "10", }, "parent_config": { "tags": [], @@ -5929,6 +8064,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "source": "loop", "step": 3, "writes": {"finish": {"my_key": " finished"}}, + "thread_id": "10", }, "parent_config": { "tags": [], @@ -5986,6 +8122,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "source": "input", "step": -1, "writes": {"__start__": {"my_key": "value", "market": "DE"}}, + "thread_id": "11", }, "parent_config": None, "next": ["__start__"], @@ -6024,6 +8161,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "source": "loop", "step": 0, "writes": None, + "thread_id": "11", }, "parent_config": { "tags": [], @@ -6095,6 +8233,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "source": "loop", "step": 1, "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "11", }, "parent_config": { "tags": [], @@ -6132,6 +8271,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "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) @@ -6155,6 +8295,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "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) @@ -6180,6 +8321,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "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) @@ -6203,6 +8345,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "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) @@ -6236,6 +8379,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "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) @@ -6259,6 +8403,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "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) @@ -6284,6 +8429,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "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) @@ -6307,6 +8453,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "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) @@ -6330,6 +8477,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "source": "update", "step": 0, "writes": {START: {"my_key": "key", "market": "DE"}}, + "thread_id": "23", }, parent_config=None, ) @@ -6352,6 +8500,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "source": "loop", "step": 1, "writes": {"prepare": {"my_key": " prepared"}}, + "thread_id": "23", }, parent_config=uconfig, ) @@ -6373,6 +8522,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "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) @@ -6724,6 +8874,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( "source": "loop", "writes": {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, "step": 4, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -7334,16 +9485,16 @@ async def test_stream_subgraphs_during_execution(checkpointer_name: str) -> None {"inner_1": {"my_key": "got here", "my_other_key": ""}}, ), ), - (FloatBetween(0.2, 0.3), ((), {"outer_1": {"my_key": " and parallel"}})), + (FloatBetween(0.2, 0.4), ((), {"outer_1": {"my_key": " and parallel"}})), ( - FloatBetween(0.5, 0.6), + FloatBetween(0.5, 0.7), ( (AnyStr("inner:"),), {"inner_2": {"my_key": " and there", "my_other_key": "got here"}}, ), ), - (FloatBetween(0.5, 0.6), ((), {"inner": {"my_key": "got here and there"}})), - (FloatBetween(0.5, 0.6), ((), {"outer_2": {"my_key": " and back again"}})), + (FloatBetween(0.5, 0.7), ((), {"inner": {"my_key": "got here and there"}})), + (FloatBetween(0.5, 0.7), ((), {"outer_2": {"my_key": " and back again"}})), ] @@ -7429,7 +9580,7 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} assert await app.ainvoke({"my_key": ""}, config, debug=True) == { - "my_key": "", + "my_key": " and parallel", } assert await app.ainvoke(None, config, debug=True) == { @@ -7464,6 +9615,7 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: c async for c in app.astream({"my_key": ""}, config, stream_mode="values") ] == [ {"my_key": ""}, + {"my_key": " and parallel"}, ] assert [c async for c in app.astream(None, config, stream_mode="values")] == [ {"my_key": ""}, @@ -7482,6 +9634,7 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: # 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")] == [ {"my_key": ""}, + {"my_key": " and parallel"}, ] assert [c async for c in app.astream(None, config, stream_mode="values")] == [ {"my_key": ""}, @@ -7496,6 +9649,7 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: c async for c in app.astream({"my_key": ""}, config, stream_mode="values") ] == [ {"my_key": ""}, + {"my_key": " and parallel"}, ] assert [c async for c in app.astream(None, config, stream_mode="values")] == [ {"my_key": ""}, @@ -7570,15 +9724,27 @@ async def test_doubly_nested_graph_interrupts(checkpointer_name: str) -> None: } # test stream updates w/ nested interrupt - config = {"configurable": {"thread_id": "2"}} + 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)] == [ {"parent_1": {"my_key": "hi my value"}}, {"__interrupt__": ()}, ] + assert nodes == ["parent_1", "grandchild_1"] assert [c async for c in app.astream(None, config)] == [ {"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"}} @@ -7677,6 +9843,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, "step": 1, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -7703,9 +9870,8 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: tasks=( PregelTask( AnyStr(), - name="inner_2", - path=(PULL, "inner_2"), - error=None, + "inner_2", + (PULL, "inner_2"), ), ), next=("inner_2",), @@ -7731,6 +9897,13 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: } }, "step": 1, + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "langgraph_node": "inner", + "langgraph_path": [PULL, "inner"], + "langgraph_step": 2, + "langgraph_triggers": ["outer_1"], + "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), parent_config={ @@ -7759,6 +9932,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, "step": 1, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -7800,6 +9974,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, "step": 1, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -7833,6 +10008,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "source": "loop", "writes": None, "step": 0, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -7866,6 +10042,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "source": "input", "writes": {"__start__": {"my_key": "my value"}}, "step": -1, + "thread_id": "1", }, created_at=AnyStr(), parent_config=None, @@ -7885,7 +10062,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "checkpoint_ns": AnyStr("inner:"), "checkpoint_id": AnyStr(), "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("inner:"): AnyStr()} + {"": AnyStr(), AnyStr("child:"): AnyStr()} ), } }, @@ -7899,6 +10076,13 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: }, "step": 1, "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "langgraph_node": "inner", + "langgraph_path": [PULL, "inner"], + "langgraph_step": 2, + "langgraph_triggers": ["outer_1"], + "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), parent_config={ @@ -7907,13 +10091,11 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "checkpoint_ns": AnyStr("inner:"), "checkpoint_id": AnyStr(), "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("inner:"): AnyStr()} + {"": AnyStr(), AnyStr("child:"): AnyStr()} ), } }, - tasks=( - PregelTask(id=AnyStr(), name="inner_2", path=(PULL, "inner_2")), - ), + tasks=(PregelTask(AnyStr(), "inner_2", (PULL, "inner_2")),), ), StateSnapshot( values={"my_key": "hi my value"}, @@ -7924,7 +10106,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "checkpoint_ns": AnyStr("inner:"), "checkpoint_id": AnyStr(), "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("inner:"): AnyStr()} + {"": AnyStr(), AnyStr("child:"): AnyStr()} ), } }, @@ -7933,6 +10115,13 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "writes": None, "step": 0, "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "langgraph_node": "inner", + "langgraph_path": [PULL, "inner"], + "langgraph_step": 2, + "langgraph_triggers": ["outer_1"], + "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), parent_config={ @@ -7941,15 +10130,15 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "checkpoint_ns": AnyStr("inner:"), "checkpoint_id": AnyStr(), "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("inner:"): AnyStr()} + {"": AnyStr(), AnyStr("child:"): AnyStr()} ), } }, tasks=( PregelTask( - id=AnyStr(), - name="inner_1", - path=(PULL, "inner_1"), + AnyStr(), + "inner_1", + (PULL, "inner_1"), result={ "my_key": "hi my value here", "my_other_key": "hi my value", @@ -7966,7 +10155,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "checkpoint_ns": AnyStr("inner:"), "checkpoint_id": AnyStr(), "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("inner:"): AnyStr()} + {"": AnyStr(), AnyStr("child:"): AnyStr()} ), } }, @@ -7975,14 +10164,21 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "writes": {"__start__": {"my_key": "hi my value"}}, "step": -1, "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("inner:"), + "langgraph_node": "inner", + "langgraph_path": [PULL, "inner"], + "langgraph_step": 2, + "langgraph_triggers": ["outer_1"], + "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), parent_config=None, tasks=( PregelTask( - id=AnyStr(), - name="__start__", - path=(PULL, "__start__"), + AnyStr(), + "__start__", + (PULL, "__start__"), result={"my_key": "hi my value"}, ), ), @@ -8010,6 +10206,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "outer_2": {"my_key": "hi my value here and there and back again"} }, "step": 3, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -8043,6 +10240,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: } }, "step": 3, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -8076,6 +10274,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "source": "loop", "writes": {"inner": {"my_key": "hi my value here and there"}}, "step": 2, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -8115,6 +10314,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "source": "loop", "writes": {"outer_1": {"my_key": "hi my value"}}, "step": 1, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -8148,6 +10348,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "source": "loop", "writes": None, "step": 0, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -8181,6 +10382,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "source": "input", "writes": {"__start__": {"my_key": "my value"}}, "step": -1, + "thread_id": "1", }, created_at=AnyStr(), parent_config=None, @@ -8286,6 +10488,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "source": "loop", "writes": {"parent_1": {"my_key": "hi my value"}}, "step": 1, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -8320,9 +10523,6 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", "checkpoint_ns": AnyStr("child:"), "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), } }, metadata={ @@ -8330,6 +10530,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "source": "loop", "writes": None, "step": 0, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -8337,9 +10538,6 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "thread_id": "1", "checkpoint_ns": AnyStr("child:"), "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), } }, ).tasks[0] @@ -8379,6 +10577,13 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "source": "loop", "writes": {"grandchild_1": {"my_key": "hi my value here"}}, "step": 1, + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [PULL, AnyStr("child_1")], + "langgraph_step": 1, + "langgraph_triggers": [AnyStr("start:child_1")], }, created_at=AnyStr(), parent_config={ @@ -8451,6 +10656,16 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: } }, "step": 1, + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [AnyStr("start:child_1")], }, created_at=AnyStr(), parent_config={ @@ -8488,6 +10703,13 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "source": "loop", "writes": None, "step": 0, + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child", + "langgraph_path": [PULL, AnyStr("child")], + "langgraph_step": 2, + "langgraph_triggers": [AnyStr("parent_1")], + "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), parent_config={ @@ -8516,6 +10738,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "source": "loop", "writes": {"parent_1": {"my_key": "hi my value"}}, "step": 1, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -8563,6 +10786,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: } }, "step": 3, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -8599,6 +10823,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: } }, "step": 3, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -8624,6 +10849,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "source": "loop", "writes": {"child": {"my_key": "hi my value here and there"}}, "step": 2, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -8667,6 +10893,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "source": "loop", "writes": {"parent_1": {"my_key": "hi my value"}}, "step": 1, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -8692,6 +10919,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "source": "loop", "writes": None, "step": 0, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -8722,6 +10950,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "source": "input", "writes": {"my_key": "my value"}, "step": -1, + "thread_id": "1", }, created_at=AnyStr(), parent_config=None, @@ -8756,6 +10985,13 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "writes": {"child_1": {"my_key": "hi my value here and there"}}, "step": 1, "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child", + "langgraph_path": [PULL, AnyStr("child")], + "langgraph_step": 2, + "langgraph_triggers": [AnyStr("parent_1")], + "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), parent_config={ @@ -8788,6 +11024,13 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "writes": None, "step": 0, "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child", + "langgraph_path": [PULL, AnyStr("child")], + "langgraph_step": 2, + "langgraph_triggers": [AnyStr("parent_1")], + "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), parent_config={ @@ -8833,6 +11076,13 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "writes": {"__start__": {"my_key": "hi my value"}}, "step": -1, "parents": {"": AnyStr()}, + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child", + "langgraph_path": [PULL, AnyStr("child")], + "langgraph_step": 2, + "langgraph_triggers": [AnyStr("parent_1")], + "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), parent_config=None, @@ -8880,6 +11130,16 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: AnyStr("child:"): AnyStr(), } ), + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [AnyStr("start:child_1")], }, created_at=AnyStr(), parent_config={ @@ -8925,6 +11185,16 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: AnyStr("child:"): AnyStr(), } ), + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [AnyStr("start:child_1")], }, created_at=AnyStr(), parent_config={ @@ -8977,6 +11247,16 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: AnyStr("child:"): AnyStr(), } ), + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [AnyStr("start:child_1")], }, created_at=AnyStr(), parent_config={ @@ -9029,6 +11309,16 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: AnyStr("child:"): AnyStr(), } ), + "thread_id": "1", + "checkpoint_ns": AnyStr("child:"), + "langgraph_checkpoint_ns": AnyStr("child:"), + "langgraph_node": "child_1", + "langgraph_path": [ + PULL, + AnyStr("child_1"), + ], + "langgraph_step": 1, + "langgraph_triggers": [AnyStr("start:child_1")], }, created_at=AnyStr(), parent_config=None, @@ -9100,7 +11390,8 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: # invoke and pause at nested interrupt assert await graph.ainvoke( - {"subjects": ["cats", "dogs"]}, config={**config, "callbacks": [tracer]} + {"subjects": ["cats", "dogs"]}, + config={**config, "callbacks": [tracer]}, ) == { "subjects": ["cats", "dogs"], "jokes": [], @@ -9109,13 +11400,36 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: # check state outer_state = await graph.aget_state(config) + + if not FF_SEND_V2: + # 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"], + } + return + assert outer_state == StateSnapshot( values={"subjects": ["cats", "dogs"], "jokes": []}, tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=("__pregel_pull", "__start__"), + error=None, + interrupts=(), + state=None, + result={"subjects": ["cats", "dogs"]}, + ), PregelTask( AnyStr(), "generate_joke", - (PUSH, 0), + (PUSH, ("__pregel_pull", "__start__"), 1, AnyStr()), state={ "configurable": { "thread_id": "1", @@ -9126,7 +11440,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: PregelTask( AnyStr(), "generate_joke", - (PUSH, 1), + (PUSH, ("__pregel_pull", "__start__"), 2, AnyStr()), state={ "configurable": { "thread_id": "1", @@ -9143,20 +11457,27 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: "checkpoint_id": AnyStr(), } }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } + metadata={ + "parents": {}, + "source": "input", + "writes": { + "__start__": { + "subjects": [ + "cats", + "dogs", + ], + } + }, + "step": -1, + "thread_id": "1", }, + created_at=AnyStr(), + parent_config=None, ) # update state of dogs joke graph await graph.aupdate_state( - outer_state.tasks[1].state, {"subject": "turtles - hohoho"} + outer_state.tasks[2].state, {"subject": "turtles - hohoho"} ) # continue past interrupt @@ -9189,7 +11510,8 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: {"jokes": ["Joke about turtles - hohoho"]}, ] }, - "step": 1, + "step": 0, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -9231,7 +11553,8 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: {"jokes": ["Joke about turtles - hohoho"]}, ] }, - "step": 1, + "step": 0, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -9243,13 +11566,22 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: }, ), StateSnapshot( - values={"subjects": ["cats", "dogs"], "jokes": []}, - next=("generate_joke", "generate_joke"), + values={"jokes": []}, + next=("__start__", "generate_joke", "generate_joke"), tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=("__pregel_pull", "__start__"), + error=None, + interrupts=(), + state=None, + result={"subjects": ["cats", "dogs"]}, + ), PregelTask( AnyStr(), "generate_joke", - (PUSH, 0), + (PUSH, ("__pregel_pull", "__start__"), 1, AnyStr()), state={ "configurable": { "thread_id": "1", @@ -9261,7 +11593,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: PregelTask( AnyStr(), "generate_joke", - (PUSH, 1), + (PUSH, ("__pregel_pull", "__start__"), 2, AnyStr()), state={ "configurable": { "thread_id": "1", @@ -9278,45 +11610,18 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: "checkpoint_id": AnyStr(), } }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"jokes": []}, - tasks=( - PregelTask( - AnyStr(), - "__start__", - (PULL, "__start__"), - result={"subjects": ["cats", "dogs"]}, - ), - ), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, metadata={ "parents": {}, "source": "input", "writes": {"__start__": {"subjects": ["cats", "dogs"]}}, "step": -1, + "thread_id": "1", }, created_at=AnyStr(), parent_config=None, ), ] - assert actual_history[1] == expected_history[1] + assert actual_history == expected_history @pytest.mark.skipif( @@ -9488,6 +11793,7 @@ async def test_weather_subgraph( "writes": {"router_node": {"route": "weather"}}, "step": 1, "parents": {}, + "thread_id": "1", }, created_at=AnyStr(), parent_config={ @@ -9579,6 +11885,7 @@ async def test_weather_subgraph( "writes": {"router_node": {"route": "weather"}}, "step": 1, "parents": {}, + "thread_id": "14", }, created_at=AnyStr(), parent_config={ @@ -9619,6 +11926,15 @@ async def test_weather_subgraph( "writes": {"model_node": {"city": "San Francisco"}}, "step": 1, "parents": {"": AnyStr()}, + "thread_id": "14", + "checkpoint_ns": AnyStr("weather_graph:"), + "langgraph_node": "weather_graph", + "langgraph_path": [PULL, "weather_graph"], + "langgraph_step": 2, + "langgraph_triggers": [ + "branch:router_node:route_after_prediction:weather_graph" + ], + "langgraph_checkpoint_ns": AnyStr("weather_graph:"), }, created_at=AnyStr(), parent_config={ @@ -9669,6 +11985,7 @@ async def test_weather_subgraph( "writes": {"router_node": {"route": "weather"}}, "step": 1, "parents": {}, + "thread_id": "14", }, created_at=AnyStr(), parent_config={ @@ -9706,8 +12023,8 @@ async def test_weather_subgraph( } }, metadata={ - "source": "update", "step": 2, + "source": "update", "writes": { "weather_node": { "messages": [ @@ -9716,6 +12033,16 @@ async def test_weather_subgraph( } }, "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:router_node:route_after_prediction:weather_graph" + ], + "langgraph_checkpoint_ns": AnyStr("weather_graph:"), }, created_at=AnyStr(), parent_config={ diff --git a/libs/langgraph/tests/test_remote_pregel.py b/libs/langgraph/tests/test_remote_graph.py similarity index 59% rename from libs/langgraph/tests/test_remote_pregel.py rename to libs/langgraph/tests/test_remote_graph.py index f2ca53583..70857ed61 100644 --- a/libs/langgraph/tests/test_remote_pregel.py +++ b/libs/langgraph/tests/test_remote_graph.py @@ -7,7 +7,9 @@ from langchain_core.runnables.graph import ( from langchain_core.runnables.graph import ( Node as DrawableNode, ) +from langgraph_sdk.schema import StreamPart +from langgraph.errors import GraphInterrupt from langgraph.pregel.remote import RemoteGraph from langgraph.pregel.types import StateSnapshot @@ -15,7 +17,7 @@ from langgraph.pregel.types import StateSnapshot def test_with_config(): # set up test remote_pregel = RemoteGraph( - graph_id="test_graph_id", + "test_graph_id", config={ "configurable": { "foo": "bar", @@ -52,7 +54,7 @@ def test_get_graph(): "type": "runnable", "data": { "id": ["langgraph", "utils", "RunnableCallable"], - "name": "agent", + "name": "agent_1", }, }, ], @@ -62,20 +64,22 @@ def test_get_graph(): ], } - remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id") + remote_pregel = RemoteGraph("test_graph_id", sync_client=mock_sync_client) # call method / assertions drawable_graph = remote_pregel.get_graph() assert drawable_graph.nodes == { "__start__": DrawableNode( - id="__start__", name="", data="__start__", metadata=None + id="__start__", name="__start__", data="__start__", metadata=None + ), + "__end__": DrawableNode( + id="__end__", name="__end__", data="__end__", metadata=None ), - "__end__": DrawableNode(id="__end__", name="", data="__end__", metadata=None), "agent": DrawableNode( id="agent", - name="", - data={"id": ["langgraph", "utils", "RunnableCallable"], "name": "agent"}, + name="agent_1", + data={"id": ["langgraph", "utils", "RunnableCallable"], "name": "agent_1"}, metadata=None, ), } @@ -99,7 +103,7 @@ async def test_aget_graph(): "type": "runnable", "data": { "id": ["langgraph", "utils", "RunnableCallable"], - "name": "agent", + "name": "agent_1", }, }, ], @@ -109,20 +113,22 @@ async def test_aget_graph(): ], } - remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id") + remote_pregel = RemoteGraph("test_graph_id", client=mock_async_client) # call method / assertions drawable_graph = await remote_pregel.aget_graph() assert drawable_graph.nodes == { "__start__": DrawableNode( - id="__start__", name="", data="__start__", metadata=None + id="__start__", name="__start__", data="__start__", metadata=None + ), + "__end__": DrawableNode( + id="__end__", name="__end__", data="__end__", metadata=None ), - "__end__": DrawableNode(id="__end__", name="", data="__end__", metadata=None), "agent": DrawableNode( id="agent", - name="", - data={"id": ["langgraph", "utils", "RunnableCallable"], "name": "agent"}, + name="agent_1", + data={"id": ["langgraph", "utils", "RunnableCallable"], "name": "agent_1"}, metadata=None, ), } @@ -133,92 +139,6 @@ async def test_aget_graph(): ] -def test_get_subgraphs(): - # set up test - mock_sync_client = MagicMock() - mock_sync_client.assistants.get_subgraphs.return_value = { - "namespace_1": { - "graph_id": "test_graph_id_2", - "input_schema": {}, - "output_schema": {}, - "state_schema": {}, - "config_schema": {}, - }, - "namespace_2": { - "graph_id": "test_graph_id_3", - "input_schema": {}, - "output_schema": {}, - "state_schema": {}, - "config_schema": {}, - }, - } - - remote_pregel = RemoteGraph( - sync_client=mock_sync_client, graph_id="test_graph_id_1" - ) - - # call method / assertions - subgraphs = list(remote_pregel.get_subgraphs()) - assert len(subgraphs) == 2 - - subgraph_1 = subgraphs[0] - ns_1 = subgraph_1[0] - remote_pregel_1: RemoteGraph = subgraph_1[1] - assert ns_1 == "namespace_1" - assert remote_pregel_1.graph_id == "test_graph_id_2" - - subgraph_2 = subgraphs[1] - ns_2 = subgraph_2[0] - remote_pregel_2: RemoteGraph = subgraph_2[1] - assert ns_2 == "namespace_2" - assert remote_pregel_2.graph_id == "test_graph_id_3" - - -@pytest.mark.anyio -async def test_aget_subgraphs(): - # set up test - mock_async_client = AsyncMock() - mock_async_client.assistants.get_subgraphs.return_value = { - "namespace_1": { - "graph_id": "test_graph_id_2", - "input_schema": {}, - "output_schema": {}, - "state_schema": {}, - "config_schema": {}, - }, - "namespace_2": { - "graph_id": "test_graph_id_3", - "input_schema": {}, - "output_schema": {}, - "state_schema": {}, - "config_schema": {}, - }, - } - - remote_pregel = RemoteGraph( - client=mock_async_client, - graph_id="test_graph_id_1", - ) - - # call method / assertions - subgraphs = [] - async for subgraph in remote_pregel.aget_subgraphs(): - subgraphs.append(subgraph) - assert len(subgraphs) == 2 - - subgraph_1 = subgraphs[0] - ns_1 = subgraph_1[0] - remote_pregel_1: RemoteGraph = subgraph_1[1] - assert ns_1 == "namespace_1" - assert remote_pregel_1.graph_id == "test_graph_id_2" - - subgraph_2 = subgraphs[1] - ns_2 = subgraph_2[0] - remote_pregel_2: RemoteGraph = subgraph_2[1] - assert ns_2 == "namespace_2" - assert remote_pregel_2.graph_id == "test_graph_id_3" - - def test_get_state(): # set up test mock_sync_client = MagicMock() @@ -238,7 +158,10 @@ def test_get_state(): } # call method / assertions - remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id") + remote_pregel = RemoteGraph( + "test_graph_id", + sync_client=mock_sync_client, + ) config = {"configurable": {"thread_id": "thread1"}} state_snapshot = remote_pregel.get_state(config) @@ -285,7 +208,10 @@ async def test_aget_state(): } # call method / assertions - remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id") + remote_pregel = RemoteGraph( + "test_graph_id", + client=mock_async_client, + ) config = {"configurable": {"thread_id": "thread1"}} state_snapshot = await remote_pregel.aget_state(config) @@ -336,7 +262,10 @@ def test_get_state_history(): ] # call method / assertions - remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id") + remote_pregel = RemoteGraph( + "test_graph_id", + sync_client=mock_sync_client, + ) config = {"configurable": {"thread_id": "thread1"}} state_history_snapshot = list( @@ -384,7 +313,10 @@ async def test_aget_state_history(): ] # call method / assertions - remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id") + remote_pregel = RemoteGraph( + "test_graph_id", + client=mock_async_client, + ) config = {"configurable": {"thread_id": "thread1"}} state_history_snapshot = [] @@ -425,7 +357,10 @@ def test_update_state(): } # call method / assertions - remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id") + remote_pregel = RemoteGraph( + "test_graph_id", + sync_client=mock_sync_client, + ) config = {"configurable": {"thread_id": "thread1"}} response = remote_pregel.update_state(config, {"key": "value"}) @@ -454,7 +389,10 @@ async def test_aupdate_state(): } # call method / assertions - remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id") + remote_pregel = RemoteGraph( + "test_graph_id", + client=mock_async_client, + ) config = {"configurable": {"thread_id": "thread1"}} response = await remote_pregel.aupdate_state(config, {"key": "value"}) @@ -473,17 +411,100 @@ def test_stream(): # set up test mock_sync_client = MagicMock() mock_sync_client.runs.stream.return_value = [ + StreamPart(event="values", data={"chunk": "data1"}), + StreamPart(event="values", data={"chunk": "data2"}), + StreamPart(event="values", data={"chunk": "data3"}), + StreamPart(event="updates", data={"chunk": "data4"}), + StreamPart(event="updates", data={"__interrupt__": ()}), + ] + + # call method / assertions + remote_pregel = RemoteGraph( + "test_graph_id", + sync_client=mock_sync_client, + ) + + # stream modes doesn't include 'updates' + stream_parts = [] + with pytest.raises(GraphInterrupt): + for stream_part in remote_pregel.stream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + stream_mode="values", + ): + stream_parts.append(stream_part) + + assert stream_parts == [ {"chunk": "data1"}, {"chunk": "data2"}, {"chunk": "data3"}, ] - # call method / assertions - remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id") + mock_sync_client.runs.stream.return_value = [ + StreamPart(event="updates", data={"chunk": "data3"}), + StreamPart(event="updates", data={"chunk": "data4"}), + StreamPart(event="updates", data={"__interrupt__": ()}), + ] - config = {"configurable": {"thread_id": "thread_1"}} - result = list(remote_pregel.stream({"input": "data"}, config)) - assert result == [{"chunk": "data1"}, {"chunk": "data2"}, {"chunk": "data3"}] + # default stream_mode is updates + stream_parts = [] + with pytest.raises(GraphInterrupt): + for stream_part in remote_pregel.stream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + {"chunk": "data3"}, + {"chunk": "data4"}, + ] + + # list stream_mode includes mode names + stream_parts = [] + with pytest.raises(GraphInterrupt): + for stream_part in remote_pregel.stream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + stream_mode=["updates"], + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + ("updates", {"chunk": "data3"}), + ("updates", {"chunk": "data4"}), + ] + + # subgraphs + list modes + stream_parts = [] + with pytest.raises(GraphInterrupt): + for stream_part in remote_pregel.stream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + stream_mode=["updates"], + subgraphs=True, + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + ((), "updates", {"chunk": "data3"}), + ((), "updates", {"chunk": "data4"}), + ] + + # subgraphs + single mode + stream_parts = [] + with pytest.raises(GraphInterrupt): + for stream_part in remote_pregel.stream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + subgraphs=True, + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + ((), {"chunk": "data3"}), + ((), {"chunk": "data4"}), + ] @pytest.mark.anyio @@ -492,57 +513,195 @@ async def test_astream(): mock_async_client = MagicMock() async_iter = MagicMock() async_iter.__aiter__.return_value = [ - {"chunk": "data1"}, - {"chunk": "data2"}, - {"chunk": "data3"}, + StreamPart(event="values", data={"chunk": "data1"}), + StreamPart(event="values", data={"chunk": "data2"}), + StreamPart(event="values", data={"chunk": "data3"}), + StreamPart(event="updates", data={"chunk": "data4"}), + StreamPart(event="updates", data={"__interrupt__": ()}), ] mock_async_client.runs.stream.return_value = async_iter # call method / assertions - remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id") + remote_pregel = RemoteGraph( + "test_graph_id", + client=mock_async_client, + ) - config = {"configurable": {"thread_id": "thread_1"}} - chunks = [] - async for chunk in remote_pregel.astream({"input": "data"}, config): - chunks.append(chunk) - assert chunks == [{"chunk": "data1"}, {"chunk": "data2"}, {"chunk": "data3"}] + # stream modes doesn't include 'updates' + stream_parts = [] + with pytest.raises(GraphInterrupt): + async for stream_part in remote_pregel.astream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + stream_mode="values", + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + {"chunk": "data1"}, + {"chunk": "data2"}, + {"chunk": "data3"}, + ] + + async_iter = MagicMock() + async_iter.__aiter__.return_value = [ + StreamPart(event="updates", data={"chunk": "data3"}), + StreamPart(event="updates", data={"chunk": "data4"}), + StreamPart(event="updates", data={"__interrupt__": ()}), + ] + mock_async_client.runs.stream.return_value = async_iter + + # default stream_mode is updates + stream_parts = [] + with pytest.raises(GraphInterrupt): + async for stream_part in remote_pregel.astream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + {"chunk": "data3"}, + {"chunk": "data4"}, + ] + + # list stream_mode includes mode names + stream_parts = [] + with pytest.raises(GraphInterrupt): + async for stream_part in remote_pregel.astream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + stream_mode=["updates"], + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + ("updates", {"chunk": "data3"}), + ("updates", {"chunk": "data4"}), + ] + + # subgraphs + list modes + stream_parts = [] + with pytest.raises(GraphInterrupt): + async for stream_part in remote_pregel.astream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + stream_mode=["updates"], + subgraphs=True, + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + ((), "updates", {"chunk": "data3"}), + ((), "updates", {"chunk": "data4"}), + ] + + # subgraphs + single mode + stream_parts = [] + with pytest.raises(GraphInterrupt): + async for stream_part in remote_pregel.astream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + subgraphs=True, + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + ((), {"chunk": "data3"}), + ((), {"chunk": "data4"}), + ] + + async_iter = MagicMock() + async_iter.__aiter__.return_value = [ + StreamPart(event="updates|my|subgraph", data={"chunk": "data3"}), + StreamPart(event="updates|hello|subgraph", data={"chunk": "data4"}), + StreamPart(event="updates|bye|subgraph", data={"__interrupt__": ()}), + ] + mock_async_client.runs.stream.return_value = async_iter + + # subgraphs + list modes + stream_parts = [] + with pytest.raises(GraphInterrupt): + async for stream_part in remote_pregel.astream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + stream_mode=["updates"], + subgraphs=True, + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + (("my", "subgraph"), "updates", {"chunk": "data3"}), + (("hello", "subgraph"), "updates", {"chunk": "data4"}), + ] + + # subgraphs + single mode + stream_parts = [] + with pytest.raises(GraphInterrupt): + async for stream_part in remote_pregel.astream( + {"input": "data"}, + config={"configurable": {"thread_id": "thread_1"}}, + subgraphs=True, + ): + stream_parts.append(stream_part) + + assert stream_parts == [ + (("my", "subgraph"), {"chunk": "data3"}), + (("hello", "subgraph"), {"chunk": "data4"}), + ] def test_invoke(): # set up test mock_sync_client = MagicMock() - mock_sync_client.runs.wait.return_value = { - "values": {"messages": [{"type": "human", "content": "world"}]} - } + mock_sync_client.runs.stream.return_value = [ + StreamPart(event="values", data={"chunk": "data1"}), + StreamPart(event="values", data={"chunk": "data2"}), + StreamPart( + event="values", data={"messages": [{"type": "human", "content": "world"}]} + ), + ] # call method / assertions - remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id") + remote_pregel = RemoteGraph( + "test_graph_id", + sync_client=mock_sync_client, + ) config = {"configurable": {"thread_id": "thread_1"}} result = remote_pregel.invoke( {"input": {"messages": [{"type": "human", "content": "hello"}]}}, config ) - assert result == {"values": {"messages": [{"type": "human", "content": "world"}]}} + assert result == {"messages": [{"type": "human", "content": "world"}]} @pytest.mark.anyio async def test_ainvoke(): # set up test - mock_async_client = AsyncMock() - mock_async_client.runs.wait.return_value = { - "values": {"messages": [{"type": "human", "content": "world"}]} - } + mock_async_client = MagicMock() + async_iter = MagicMock() + async_iter.__aiter__.return_value = [ + StreamPart(event="values", data={"chunk": "data1"}), + StreamPart(event="values", data={"chunk": "data2"}), + StreamPart( + event="values", data={"messages": [{"type": "human", "content": "world"}]} + ), + ] + mock_async_client.runs.stream.return_value = async_iter # call method / assertions - remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id") + remote_pregel = RemoteGraph( + "test_graph_id", + client=mock_async_client, + ) config = {"configurable": {"thread_id": "thread_1"}} result = await remote_pregel.ainvoke( {"input": {"messages": [{"type": "human", "content": "hello"}]}}, config ) - assert result == {"values": {"messages": [{"type": "human", "content": "world"}]}} + assert result == {"messages": [{"type": "human", "content": "world"}]} @pytest.mark.skip("Unskip this test to manually test the LangGraph Cloud integration") @@ -557,7 +716,9 @@ async def test_langgraph_cloud_integration(): client = get_client() sync_client = get_sync_client() remote_pregel = RemoteGraph( - client=client, sync_client=sync_client, graph_id="agent" + "agent", + client=client, + sync_client=sync_client, ) # define graph @@ -572,7 +733,7 @@ async def test_langgraph_cloud_integration(): "messages": [ { "role": "human", - "content": "Hello world!", + "content": "What's the weather in SF?", } ] } @@ -580,7 +741,8 @@ async def test_langgraph_cloud_integration(): # test invoke response = app.invoke( input, - config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}}, + config={"configurable": {"thread_id": "39a6104a-34e7-4f83-929c-d9eb163003c9"}}, + interrupt_before=["agent"], ) print("response:", response["messages"][-1].content) @@ -634,9 +796,3 @@ async def test_langgraph_cloud_integration(): remote_pregel.graph_id = "fe096781-5601-53d2-b2f6-0d3403f7e9ca" # must be UUID graph = await remote_pregel.aget_graph(xray=True) print("graph:", graph) - - # test get subgraphs - remote_pregel.graph_id = "fe096781-5601-53d2-b2f6-0d3403f7e9ca" # must be UUID - async for name, pregel in remote_pregel.aget_subgraphs(): - print("name:", name) - print("pregel:", pregel) diff --git a/libs/langgraph/tests/test_state.py b/libs/langgraph/tests/test_state.py index 73546b6a1..0a4a8725a 100644 --- a/libs/langgraph/tests/test_state.py +++ b/libs/langgraph/tests/test_state.py @@ -5,11 +5,11 @@ from typing import Annotated as Annotated2 from typing import Any, Optional import pytest -from langchain_core.runnables import RunnableConfig +from langchain_core.runnables import RunnableConfig, RunnableLambda from pydantic.v1 import BaseModel from typing_extensions import Annotated, NotRequired, Required, TypedDict -from langgraph.graph.state import StateGraph, _warn_invalid_state_schema +from langgraph.graph.state import StateGraph, _get_node_name, _warn_invalid_state_schema from langgraph.managed.shared_value import SharedValue @@ -61,6 +61,9 @@ def test_state_schema_with_type_hint(): class OutputState(TypedDict): input_state: InputState + class FooState(InputState): + foo: str + def complete_hint(state: InputState) -> OutputState: return {"input_state": state} @@ -73,24 +76,46 @@ def test_state_schema_with_type_hint(): def miss_all_hint(state, config): return {"input_state": state} + def pre_foo(_) -> FooState: + return {"foo": "bar"} + + class Foo: + def __call__(self, state: FooState) -> OutputState: + assert state.pop("foo") == "bar" + return {"input_state": state} + graph = StateGraph(InputState, output=OutputState) - actions = [complete_hint, miss_first_hint, only_return_hint, miss_all_hint] + actions = [ + complete_hint, + miss_first_hint, + only_return_hint, + miss_all_hint, + pre_foo, + Foo(), + ] for action in actions: graph.add_node(action) - graph.set_entry_point(actions[0].__name__) + def get_name(action) -> str: + return getattr(action, "__name__", action.__class__.__name__) + + graph.set_entry_point(get_name(actions[0])) for i in range(len(actions) - 1): - graph.add_edge(actions[i].__name__, actions[i + 1].__name__) - graph.set_finish_point(actions[-1].__name__) + graph.add_edge(get_name(actions[i]), get_name(actions[i + 1])) + graph.set_finish_point(get_name(actions[-1])) graph = graph.compile() input_state = InputState(question="Hello World!") output_state = OutputState(input_state=input_state) + foo_state = FooState(foo="bar") for i, c in enumerate(graph.stream(input_state, stream_mode="updates")): - node_name = actions[i].__name__ - assert c[node_name] == output_state + node_name = get_name(actions[i]) + if node_name == get_name(pre_foo): + assert c[node_name] == foo_state + else: + assert c[node_name] == output_state @pytest.mark.parametrize("total_", [True, False]) @@ -261,3 +286,35 @@ def test_raises_invalid_managed(): match="Invalid managed channels detected in BadOutputState: some_output_channel. Managed channels are not permitted in Input/Output schema.", ): StateGraph(_state, input=_inp, output=_outp) + + +def test__get_node_name() -> None: + # default runnable name + assert _get_node_name(RunnableLambda(func=lambda x: x)) == "RunnableLambda" + # custom runnable name + assert ( + _get_node_name(RunnableLambda(name="my_runnable", func=lambda x: x)) + == "my_runnable" + ) + + # lambda + assert _get_node_name(lambda x: x) == "