diff --git a/.github/workflows/_test_langgraph.yml b/.github/workflows/_test_langgraph.yml index b8e2679fe..b771da867 100644 --- a/.github/workflows/_test_langgraph.yml +++ b/.github/workflows/_test_langgraph.yml @@ -19,9 +19,13 @@ jobs: - "3.13" core-version: - "latest" + ff-send-v2: + - "false" include: - python-version: "3.11" core-version: ">=0.2.42,<0.3.0" + - python-version: "3.11" + ff-send-v2: "true" defaults: run: @@ -52,6 +56,8 @@ jobs: - name: Run tests shell: bash + env: + LANGGRAPH_FF_SEND_V2: ${{ matrix.ff-send-v2 }} run: | make test diff --git a/libs/langgraph/Makefile b/libs/langgraph/Makefile index 1e249a0cd..2aacf6db8 100644 --- a/libs/langgraph/Makefile +++ b/libs/langgraph/Makefile @@ -49,7 +49,7 @@ test: exit $$EXIT_CODE test_watch: - make start-postgres && poetry run ptw . -- --ff -v -x -n auto --dist worksteal --snapshot-update --tb short $(TEST); \ + make start-postgres && poetry run ptw . -- --ff -vv -x -n auto --dist worksteal --snapshot-update --tb short $(TEST); \ EXIT_CODE=$$?; \ make stop-postgres; \ exit $$EXIT_CODE diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index 5d18b5262..c987dd69a 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -1,4 +1,5 @@ import sys +from os import getenv from types import MappingProxyType from typing import Any, Literal, Mapping, cast @@ -81,6 +82,8 @@ NS_END = sys.intern(":") # for checkpoint_ns, for each level, separates the namespace from the task_id CONF = cast(Literal["configurable"], sys.intern("configurable")) # key for the configurable dict in RunnableConfig +FF_SEND_V2 = getenv("LANGGRAPH_FF_SEND_V2", "false").lower() == "true" +# temporary flag to enable new Send semantics RESERVED = { TAG_HIDDEN, diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 6ae0253ac..d6669bdb5 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -173,7 +173,7 @@ def local_write( """Function injected under CONFIG_KEY_SEND in task config, to write to channels. Validates writes and forwards them to `commit` function.""" for chan, value in writes: - if chan == PUSH: + if chan in (PUSH, TASKS): if not isinstance(value, Send): raise InvalidUpdateError(f"Expected Send, got {value}") if value.node not in process_keys: diff --git a/libs/langgraph/langgraph/pregel/write.py b/libs/langgraph/langgraph/pregel/write.py index ba783453c..3af0fe5e9 100644 --- a/libs/langgraph/langgraph/pregel/write.py +++ b/libs/langgraph/langgraph/pregel/write.py @@ -14,7 +14,7 @@ from typing import ( from langchain_core.runnables import Runnable, RunnableConfig from langchain_core.runnables.utils import ConfigurableFieldSpec -from langgraph.constants import CONF, CONFIG_KEY_SEND, PUSH, TASKS, Send +from langgraph.constants import CONF, CONFIG_KEY_SEND, FF_SEND_V2, PUSH, TASKS, Send from langgraph.errors import InvalidUpdateError from langgraph.utils.runnable import RunnableCallable @@ -119,7 +119,11 @@ class ChannelWrite(RunnableCallable): if w.value is PASSTHROUGH: raise InvalidUpdateError("PASSTHROUGH value must be replaced") # split packets and entries - sends = [(PUSH, packet) for packet in writes if isinstance(packet, Send)] + sends = [ + (PUSH if FF_SEND_V2 else TASKS, packet) + for packet in writes + if isinstance(packet, Send) + ] entries = [write for write in writes if isinstance(write, ChannelWriteEntry)] # process entries into values values = [ diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index f22f96de5..eae7694ff 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -327,7 +327,6 @@ async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]: ALL_CHECKPOINTERS_SYNC = [ "memory", "sqlite", - "duckdb", "postgres", "postgres_pipe", "postgres_pool", diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index c57948c00..bce736ac5 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -54,7 +54,14 @@ from langgraph.checkpoint.base import ( CheckpointTuple, ) from langgraph.checkpoint.memory import MemorySaver -from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH, START +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, GraphCommand, StateGraph from langgraph.graph.message import MessageGraph, MessagesState, add_messages @@ -1781,17 +1788,31 @@ def test_concurrent_emit_sends() -> None: 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", - ] + 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: @@ -1831,16 +1852,31 @@ def test_send_sequences() -> None: 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", - ] + 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) @@ -1848,8 +1884,8 @@ def test_send_sequences() -> None: def test_send_dedupe_on_resume( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: - if checkpointer_name == "duckdb": - pytest.skip("DuckDB isn't returning the right history") + if not FF_SEND_V2: + pytest.skip("Send deduplication is only available in Send V2") checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") class InterruptOnce: @@ -2298,6 +2334,9 @@ def test_send_react_interrupt( } 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( @@ -2741,6 +2780,9 @@ def test_send_react_interrupt_control( } 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"]) @@ -5866,6 +5908,9 @@ def test_state_graph_packets( {"__interrupt__": ()}, ] + if not FF_SEND_V2: + return + assert app_w_interrupt.get_state(config) == StateSnapshot( values={ "messages": [ @@ -12264,6 +12309,21 @@ 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=( @@ -12409,7 +12469,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[2].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( diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index cc1bd0d39..c82935082 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -51,7 +51,14 @@ from langgraph.checkpoint.base import ( CheckpointTuple, ) from langgraph.checkpoint.memory import MemorySaver -from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH, START +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, GraphCommand, StateGraph from langgraph.graph.message import MessageGraph, MessagesState, add_messages @@ -2022,17 +2029,31 @@ async def test_concurrent_emit_sends() -> None: 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", - ] + 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) @@ -2073,16 +2094,34 @@ async def test_send_sequences(checkpointer_name: str) -> None: 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", - ] + 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"]) @@ -2110,6 +2149,9 @@ async def test_send_sequences(checkpointer_name: str) -> None: @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 @@ -2542,6 +2584,9 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: } 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( @@ -3004,6 +3049,9 @@ async def test_send_react_interrupt_control(checkpointer_name: str) -> None: } 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( @@ -5826,6 +5874,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": [ @@ -11029,6 +11080,20 @@ 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=(