diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index ba028d7bd..027f51ee7 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -1,3 +1,4 @@ +import binascii import itertools import sys from collections import defaultdict, deque @@ -18,7 +19,6 @@ from typing import ( cast, overload, ) -from uuid import UUID from langchain_core.callbacks import Callbacks from langchain_core.callbacks.manager import AsyncParentRunManager, ParentRunManager @@ -27,6 +27,7 @@ from langchain_core.runnables.config import RunnableConfig from langgraph.channels.base import BaseChannel from langgraph.checkpoint.base import ( BaseCheckpointSaver, + ChannelVersions, Checkpoint, PendingWrite, V, @@ -373,6 +374,8 @@ def prepare_next_tasks( """Prepare the set of tasks that will make up the next Pregel step. This is the union of all PUSH tasks (Sends) and PULL tasks (nodes triggered by edges).""" + checkpoint_id_bytes = binascii.unhexlify(checkpoint["id"].replace("-", "")) + null_version = checkpoint_null_version(checkpoint) tasks: list[Union[PregelTask, PregelExecutableTask]] = [] # Consume pending_sends from previous step for idx, _ in enumerate(checkpoint["pending_sends"]): @@ -380,6 +383,8 @@ def prepare_next_tasks( (PUSH, idx), None, checkpoint=checkpoint, + checkpoint_id_bytes=checkpoint_id_bytes, + checkpoint_null_version=null_version, pending_writes=pending_writes, processes=processes, channels=channels, @@ -399,6 +404,8 @@ def prepare_next_tasks( (PULL, name), None, checkpoint=checkpoint, + checkpoint_id_bytes=checkpoint_id_bytes, + checkpoint_null_version=null_version, pending_writes=pending_writes, processes=processes, channels=channels, @@ -414,11 +421,16 @@ def prepare_next_tasks( return {t.id: t for t in tasks} +PUSH_TRIGGER = (PUSH,) + + def prepare_single_task( task_path: tuple[Any, ...], task_id_checksum: Optional[str], *, checkpoint: Checkpoint, + checkpoint_id_bytes: bytes, + checkpoint_null_version: Optional[V], pending_writes: list[PendingWrite], processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], @@ -432,7 +444,6 @@ def prepare_single_task( ) -> Union[None, PregelTask, PregelExecutableTask]: """Prepares a single task for the next Pregel step, given a task path, which uniquely identifies a PUSH or PULL task within the graph.""" - checkpoint_id = UUID(checkpoint["id"]).bytes configurable = config.get(CONF, {}) parent_ns = configurable.get(CONFIG_KEY_CHECKPOINT_NS, "") @@ -445,10 +456,10 @@ def prepare_single_task( if name is None: raise ValueError("`call` functions must have a `__name__` attribute") # create task id - triggers = [PUSH] + triggers: Sequence[str] = PUSH_TRIGGER checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name task_id = _uuid5_str( - checkpoint_id, + checkpoint_id_bytes, checkpoint_ns, str(step), name, @@ -539,12 +550,12 @@ def prepare_single_task( ) return # create task id - triggers = [PUSH] + triggers = PUSH_TRIGGER checkpoint_ns = ( f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node ) task_id = _uuid5_str( - checkpoint_id, + checkpoint_id_bytes, checkpoint_ns, str(step), packet.node, @@ -641,18 +652,15 @@ def prepare_single_task( if name not in processes: return proc = processes[name] - version_type = type(next(iter(checkpoint["channel_versions"].values()), None)) - null_version = version_type() # type: ignore[misc] - if null_version is None: + if checkpoint_null_version is None: return - seen = checkpoint["versions_seen"].get(name, {}) # If any of the channels read by this process were updated - if triggers := sorted( - chan - for chan in proc.triggers - if channels[chan].is_available() - and checkpoint["channel_versions"].get(chan, null_version) # type: ignore[operator] - > seen.get(chan, null_version) + if triggers := _triggers( + channels, + checkpoint["channel_versions"], + checkpoint["versions_seen"].get(name), + checkpoint_null_version, + proc, ): try: val = next( @@ -670,7 +678,7 @@ def prepare_single_task( # create task id checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name task_id = _uuid5_str( - checkpoint_id, + checkpoint_id_bytes, checkpoint_ns, str(step), name, @@ -761,6 +769,35 @@ def prepare_single_task( return PregelTask(task_id, name, task_path[:3]) +def checkpoint_null_version( + checkpoint: Checkpoint, +) -> Optional[V]: + """Get the null version for the checkpoint, if available.""" + for version in checkpoint["channel_versions"].values(): + return type(version)() + return None + + +def _triggers( + channels: Mapping[str, BaseChannel], + versions: ChannelVersions, + seen: Optional[ChannelVersions], + null_version: V, + proc: PregelNode, +) -> Sequence[str]: + if seen is None: + for chan in proc.triggers: + if channels[chan].is_available(): + return (chan,) + else: + for chan in proc.triggers: + if channels[chan].is_available() and versions.get( # type: ignore[operator] + chan, null_version + ) > seen.get(chan, null_version): + return (chan,) + return EMPTY_SEQ + + def _scratchpad( config: RunnableConfig, pending_writes: list[PendingWrite], diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 7ab3431b7..edd69db01 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -1,4 +1,5 @@ import asyncio +import binascii import concurrent.futures from collections import defaultdict, deque from contextlib import AsyncExitStack, ExitStack @@ -79,6 +80,7 @@ from langgraph.pregel.algo import ( GetNextVersion, PregelTaskWrites, apply_writes, + checkpoint_null_version, increment, prepare_next_tasks, prepare_single_task, @@ -347,12 +349,16 @@ class PregelLoop(LoopProtocol): ): self.to_interrupt.append(task) return + checkpoint_id_bytes = binascii.unhexlify(self.checkpoint["id"].replace("-", "")) + null_version = checkpoint_null_version(self.checkpoint) if pushed := cast( Optional[PregelExecutableTask], prepare_single_task( (PUSH, task.path, write_idx, task.id, call), None, checkpoint=self.checkpoint, + checkpoint_id_bytes=checkpoint_id_bytes, + checkpoint_null_version=null_version, pending_writes=self.checkpoint_pending_writes, processes=self.nodes, channels=self.channels, diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 5efb89385..4c8cf6da4 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -149,7 +149,7 @@ class PregelExecutableTask(NamedTuple): proc: Runnable writes: deque[tuple[str, Any]] config: RunnableConfig - triggers: list[str] + triggers: Sequence[str] retry_policy: Optional[RetryPolicy] cache_policy: Optional[CachePolicy] id: str diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index 57c9b4071..d6637f000 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -2483,7 +2483,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: { "langgraph_step": 1, "langgraph_node": "agent", - "langgraph_triggers": ["start:agent"], + "langgraph_triggers": ("start:agent",), "langgraph_path": (PULL, "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), @@ -2500,7 +2500,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: { "langgraph_step": 2, "langgraph_node": "tools", - "langgraph_triggers": ["branch:to:tools"], + "langgraph_triggers": ("branch:to:tools",), "langgraph_path": (PULL, "tools"), "langgraph_checkpoint_ns": AnyStr("tools:"), }, @@ -2542,7 +2542,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: { "langgraph_step": 3, "langgraph_node": "agent", - "langgraph_triggers": ["tools"], + "langgraph_triggers": ("tools",), "langgraph_path": (PULL, "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), @@ -2559,7 +2559,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: { "langgraph_step": 4, "langgraph_node": "tools", - "langgraph_triggers": ["branch:to:tools"], + "langgraph_triggers": ("branch:to:tools",), "langgraph_path": (PULL, "tools"), "langgraph_checkpoint_ns": AnyStr("tools:"), }, @@ -2573,7 +2573,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: { "langgraph_step": 4, "langgraph_node": "tools", - "langgraph_triggers": ["branch:to:tools"], + "langgraph_triggers": ("branch:to:tools",), "langgraph_path": (PULL, "tools"), "langgraph_checkpoint_ns": AnyStr("tools:"), }, @@ -2585,7 +2585,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: { "langgraph_step": 5, "langgraph_node": "agent", - "langgraph_triggers": ["tools"], + "langgraph_triggers": ("tools",), "langgraph_path": (PULL, "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), @@ -5501,7 +5501,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "id": AnyStr(), "name": "rewrite_query", "input": {"query": "what is weather in sf", "docs": []}, - "triggers": ["start:rewrite_query"], + "triggers": ("start:rewrite_query",), }, }, ), @@ -5532,7 +5532,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "id": AnyStr(), "name": "retriever_one", "input": {"query": "query: what is weather in sf", "docs": []}, - "triggers": ["rewrite_query"], + "triggers": ("rewrite_query",), }, }, ), @@ -5546,7 +5546,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "id": AnyStr(), "name": "retriever_two", "input": {"query": "query: what is weather in sf", "docs": []}, - "triggers": ["rewrite_query"], + "triggers": ("rewrite_query",), }, }, ), @@ -5608,7 +5608,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None: "query": "query: what is weather in sf", "docs": ["doc1", "doc2", "doc3", "doc4"], }, - "triggers": ["retriever_one", "retriever_two"], + "triggers": (AnyStr("retriever_"),), }, }, ), @@ -6634,7 +6634,7 @@ def test_branch_then( "id": AnyStr(), "name": "prepare", "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], + "triggers": ("start:prepare",), }, }, { @@ -6706,7 +6706,7 @@ def test_branch_then( "id": AnyStr(), "name": "tool_two_slow", "input": {"my_key": "value prepared", "market": "DE"}, - "triggers": ["branch:to:tool_two_slow"], + "triggers": ("branch:to:tool_two_slow",), }, }, { @@ -6773,7 +6773,7 @@ def test_branch_then( "id": AnyStr(), "name": "finish", "input": {"my_key": "value prepared slow", "market": "DE"}, - "triggers": ["branch:prepare:condition::then"], + "triggers": ("branch:prepare:condition::then",), }, }, { diff --git a/libs/langgraph/tests/test_large_cases_async.py b/libs/langgraph/tests/test_large_cases_async.py index 95fcb9646..157e0e080 100644 --- a/libs/langgraph/tests/test_large_cases_async.py +++ b/libs/langgraph/tests/test_large_cases_async.py @@ -2300,7 +2300,7 @@ async def test_prebuilt_tool_chat() -> None: { "langgraph_step": 1, "langgraph_node": "agent", - "langgraph_triggers": ["start:agent"], + "langgraph_triggers": ("start:agent",), "langgraph_path": ("__pregel_pull", "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), @@ -2317,7 +2317,7 @@ async def test_prebuilt_tool_chat() -> None: { "langgraph_step": 2, "langgraph_node": "tools", - "langgraph_triggers": ["branch:to:tools"], + "langgraph_triggers": ("branch:to:tools",), "langgraph_path": ("__pregel_pull", "tools"), "langgraph_checkpoint_ns": AnyStr("tools:"), }, @@ -2359,7 +2359,7 @@ async def test_prebuilt_tool_chat() -> None: { "langgraph_step": 3, "langgraph_node": "agent", - "langgraph_triggers": ["tools"], + "langgraph_triggers": ("tools",), "langgraph_path": ("__pregel_pull", "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), @@ -2376,7 +2376,7 @@ async def test_prebuilt_tool_chat() -> None: { "langgraph_step": 4, "langgraph_node": "tools", - "langgraph_triggers": ["branch:to:tools"], + "langgraph_triggers": ("branch:to:tools",), "langgraph_path": ("__pregel_pull", "tools"), "langgraph_checkpoint_ns": AnyStr("tools:"), }, @@ -2390,7 +2390,7 @@ async def test_prebuilt_tool_chat() -> None: { "langgraph_step": 4, "langgraph_node": "tools", - "langgraph_triggers": ["branch:to:tools"], + "langgraph_triggers": ("branch:to:tools",), "langgraph_path": ("__pregel_pull", "tools"), "langgraph_checkpoint_ns": AnyStr("tools:"), }, @@ -2402,7 +2402,7 @@ async def test_prebuilt_tool_chat() -> None: { "langgraph_step": 5, "langgraph_node": "agent", - "langgraph_triggers": ["tools"], + "langgraph_triggers": ("tools",), "langgraph_path": ("__pregel_pull", "agent"), "langgraph_checkpoint_ns": AnyStr("agent:"), "checkpoint_ns": AnyStr("agent:"), @@ -3883,7 +3883,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "id": AnyStr(), "name": "rewrite_query", "input": {"query": "what is weather in sf", "docs": []}, - "triggers": ["start:rewrite_query"], + "triggers": ("start:rewrite_query",), }, }, ), @@ -3914,7 +3914,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "id": AnyStr(), "name": "retriever_one", "input": {"query": "query: what is weather in sf", "docs": []}, - "triggers": ["rewrite_query"], + "triggers": ("rewrite_query",), }, }, ), @@ -3928,7 +3928,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "id": AnyStr(), "name": "retriever_two", "input": {"query": "query: what is weather in sf", "docs": []}, - "triggers": ["rewrite_query"], + "triggers": ("rewrite_query",), }, }, ), @@ -3990,7 +3990,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: "query": "query: what is weather in sf", "docs": ["doc1", "doc2", "doc3", "doc4"], }, - "triggers": ["retriever_one", "retriever_two"], + "triggers": (AnyStr("retriever_"),), }, }, ), @@ -4465,7 +4465,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "id": AnyStr(), "name": "prepare", "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], + "triggers": ("start:prepare",), }, }, { @@ -4537,7 +4537,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "id": AnyStr(), "name": "tool_two_slow", "input": {"my_key": "value prepared", "market": "DE"}, - "triggers": ["branch:to:tool_two_slow"], + "triggers": ("branch:to:tool_two_slow",), }, }, { @@ -4609,7 +4609,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "id": AnyStr(), "name": "finish", "input": {"my_key": "value prepared slow", "market": "DE"}, - "triggers": ["branch:prepare:condition::then"], + "triggers": ("branch:prepare:condition::then",), }, }, { @@ -4778,7 +4778,7 @@ async def test_branch_then(checkpointer_name: str) -> None: "id": AnyStr(), "name": "prepare", "input": {"my_key": "value", "market": "DE"}, - "triggers": ["start:prepare"], + "triggers": ("start:prepare",), }, }, { diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index ccd899ce9..d8d51a8a0 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -817,7 +817,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "id": AnyStr(), "name": "one", "input": 2, - "triggers": ["input"], + "triggers": ("input",), }, }, { @@ -828,7 +828,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "id": AnyStr(), "name": "two", "input": [12], - "triggers": ["inbox"], + "triggers": ("inbox",), }, }, { @@ -863,7 +863,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "id": AnyStr(), "name": "two", "input": [3], - "triggers": ["inbox"], + "triggers": ("inbox",), }, }, { @@ -3247,14 +3247,24 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( assert [ c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config) - ] == [ - {"rewrite_query": {"query": "query: what is weather in sf"}}, - {"qa": {"answer": ""}}, - {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, - {"retriever_two": {"docs": ["doc3", "doc4"]}}, - {"retriever_one": {"docs": ["doc1", "doc2"]}}, - {"__interrupt__": ()}, - ] + ] in ( + [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"qa": {"answer": ""}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ], + [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"qa": {"answer": ""}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + {"__interrupt__": ()}, + ], + ) assert [c for c in app_w_interrupt.stream(None, config)] == [ {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, @@ -5969,9 +5979,7 @@ def test_falsy_return_from_task( "a": 5, }, "name": "graph", - "triggers": [ - "__start__", - ], + "triggers": ("__start__",), }, "step": 0, "timestamp": AnyStr(), @@ -5985,9 +5993,7 @@ def test_falsy_return_from_task( {}, ), "name": "falsy_task", - "triggers": [ - "__pregel_push", - ], + "triggers": ("__pregel_push",), }, "step": 0, "timestamp": AnyStr(), @@ -6094,9 +6100,7 @@ def test_falsy_return_from_task( "a": 5, }, "name": "graph", - "triggers": [ - "__start__", - ], + "triggers": ("__start__",), }, "step": 0, "timestamp": AnyStr(), @@ -6110,9 +6114,7 @@ def test_falsy_return_from_task( {}, ), "name": "falsy_task", - "triggers": [ - "__pregel_push", - ], + "triggers": ("__pregel_push",), }, "step": 0, "timestamp": AnyStr(), @@ -6923,7 +6925,7 @@ def test_tags_stream_mode_messages() -> None: { "langgraph_step": 1, "langgraph_node": "call_model", - "langgraph_triggers": ["start:call_model"], + "langgraph_triggers": ("start:call_model",), "langgraph_path": ("__pregel_pull", "call_model"), "langgraph_checkpoint_ns": AnyStr("call_model:"), "checkpoint_ns": AnyStr("call_model:"), diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index a42a84a28..f1bb4d8ab 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -1672,7 +1672,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "id": AnyStr(), "name": "one", "input": 2, - "triggers": ["input"], + "triggers": ("input",), }, }, { @@ -1683,7 +1683,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "id": AnyStr(), "name": "two", "input": [12], - "triggers": ["inbox"], + "triggers": ("inbox",), }, }, { @@ -1718,7 +1718,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: "id": AnyStr(), "name": "two", "input": [3], - "triggers": ["inbox"], + "triggers": ("inbox",), }, }, { @@ -7571,7 +7571,7 @@ async def test_tags_stream_mode_messages() -> None: { "langgraph_step": 1, "langgraph_node": "call_model", - "langgraph_triggers": ["start:call_model"], + "langgraph_triggers": ("start:call_model",), "langgraph_path": ("__pregel_pull", "call_model"), "langgraph_checkpoint_ns": AnyStr("call_model:"), "checkpoint_ns": AnyStr("call_model:"), diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py index fa9a221d0..2d442e131 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py @@ -1,4 +1,5 @@ import asyncio +import binascii import concurrent.futures from collections.abc import Sequence from contextlib import ( @@ -19,7 +20,7 @@ import langgraph.scheduler.kafka.serde as serde from langgraph.constants import CONFIG_KEY_DELEGATE, ERROR from langgraph.errors import CheckpointNotLatest, GraphDelegate, TaskNotFound from langgraph.pregel import Pregel -from langgraph.pregel.algo import prepare_single_task +from langgraph.pregel.algo import checkpoint_null_version, prepare_single_task from langgraph.pregel.executor import ( AsyncBackgroundExecutor, BackgroundExecutor, @@ -209,6 +210,10 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager): for_execution=True, checkpointer=self.graph.checkpointer, store=self.graph.store, + checkpoint_id_bytes=binascii.unhexlify( + saved.checkpoint["id"].replace("-", "") + ), + checkpoint_null_version=checkpoint_null_version(saved.checkpoint), ): # execute task, saving writes runner = PregelRunner( @@ -421,6 +426,10 @@ class KafkaExecutor(AbstractContextManager): step=saved.metadata["step"] + 1, for_execution=True, checkpointer=self.graph.checkpointer, + checkpoint_id_bytes=binascii.unhexlify( + saved.checkpoint["id"].replace("-", "") + ), + checkpoint_null_version=checkpoint_null_version(saved.checkpoint), ): # execute task, saving writes runner = PregelRunner(