diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 4031516c8..12a53b8ad 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -236,7 +236,7 @@ def apply_writes( # sort tasks on path, to ensure deterministic order for update application # any path parts after the 3rd are ignored for sorting # (we use them for eg. task ids which aren't good for sorting) - tasks = sorted(tasks, key=lambda t: _tuple_str(t.path[:3])) + tasks = sorted(tasks, key=lambda t: task_path_str(t.path[:3])) # if no task has triggers this is applying writes from the null task only # so we don't do anything other than update the channels written to bump_step = any(t.triggers for t in tasks) @@ -450,7 +450,7 @@ def prepare_single_task( str(step), name, PUSH, - _tuple_str(task_path[1]), + task_path_str(task_path[1]), str(task_path[2]), ) task_checkpoint_ns = f"{checkpoint_ns}:{task_id}" @@ -813,10 +813,10 @@ def _uuid5_str(namespace: bytes, *parts: str) -> str: return f"{hex[:8]}-{hex[8:12]}-{hex[12:16]}-{hex[16:20]}-{hex[20:32]}" -def _tuple_str(tup: Union[str, int, tuple]) -> str: - """Generate a string representation of a tuple.""" +def task_path_str(tup: Union[str, int, tuple]) -> str: + """Generate a string representation of the task path.""" return ( - f"~{', '.join(_tuple_str(x) for x in tup)}" + f"~{', '.join(task_path_str(x) for x in tup)}" if isinstance(tup, (tuple, list)) else f"{tup:010d}" if isinstance(tup, int) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 632219ac9..5fd9b5d96 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -2,6 +2,7 @@ import asyncio import concurrent.futures from collections import defaultdict, deque from contextlib import AsyncExitStack, ExitStack +from inspect import signature from types import TracebackType from typing import ( Any, @@ -81,6 +82,7 @@ from langgraph.pregel.algo import ( prepare_next_tasks, prepare_single_task, should_interrupt, + task_path_str, ) from langgraph.pregel.debug import ( map_debug_checkpoint, @@ -151,6 +153,7 @@ class PregelLoop(LoopProtocol): checkpointer_put_writes: Optional[ Callable[[RunnableConfig, Sequence[tuple[str, Any]], str], Any] ] + checkpointer_put_writes_accepts_task_path: bool _checkpointer_put_after_previous: Optional[ Callable[ [ @@ -288,20 +291,34 @@ class PregelLoop(LoopProtocol): else: self.checkpoint_pending_writes.append((task_id, c, v)) if self.checkpointer_put_writes is not None: - self.submit( - self.checkpointer_put_writes, - patch_configurable( - self.checkpoint_config, - { - CONFIG_KEY_CHECKPOINT_NS: self.config[CONF].get( - CONFIG_KEY_CHECKPOINT_NS, "" - ), - CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"], - }, - ), - writes, - task_id, + config = patch_configurable( + self.checkpoint_config, + { + CONFIG_KEY_CHECKPOINT_NS: self.config[CONF].get( + CONFIG_KEY_CHECKPOINT_NS, "" + ), + CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"], + }, ) + if self.checkpointer_put_writes_accepts_task_path: + if hasattr(self, "tasks"): + task = self.tasks.get(task_id) + else: + task = None + self.submit( + self.checkpointer_put_writes, + config, + writes, + task_id, + task_path_str(task.path) if task else "", + ) + else: + self.submit( + self.checkpointer_put_writes, + config, + writes, + task_id, + ) # output writes if hasattr(self, "tasks"): self._output_writes(task_id, writes) @@ -813,10 +830,15 @@ class SyncPregelLoop(PregelLoop, ContextManager): if checkpointer: self.checkpointer_get_next_version = checkpointer.get_next_version self.checkpointer_put_writes = checkpointer.put_writes + self.checkpointer_put_writes_accepts_task_path = ( + signature(checkpointer.put_writes).parameters.get("task_path") + is not None + ) else: self.checkpointer_get_next_version = increment self._checkpointer_put_after_previous = None # type: ignore[assignment] self.checkpointer_put_writes = None + self.checkpointer_put_writes_accepts_task_path = False def _checkpointer_put_after_previous( self, @@ -945,10 +967,15 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): if checkpointer: self.checkpointer_get_next_version = checkpointer.get_next_version self.checkpointer_put_writes = checkpointer.aput_writes + self.checkpointer_put_writes_accepts_task_path = ( + signature(checkpointer.aput_writes).parameters.get("task_path") + is not None + ) else: self.checkpointer_get_next_version = increment self._checkpointer_put_after_previous = None # type: ignore[assignment] self.checkpointer_put_writes = None + self.checkpointer_put_writes_accepts_task_path = False async def _checkpointer_put_after_previous( self, diff --git a/libs/langgraph/poetry.lock b/libs/langgraph/poetry.lock index 7a9e9eb79..9e56cd046 100644 --- a/libs/langgraph/poetry.lock +++ b/libs/langgraph/poetry.lock @@ -1348,7 +1348,7 @@ typing-extensions = ">=4.7" [[package]] name = "langgraph-checkpoint" -version = "2.0.9" +version = "2.0.10" description = "Library with base interfaces for LangGraph checkpoint savers." optional = false python-versions = "^3.9.0,<4.0" @@ -1366,7 +1366,7 @@ url = "../checkpoint" [[package]] name = "langgraph-checkpoint-postgres" -version = "2.0.11" +version = "2.0.12" description = "Library with a Postgres implementation of LangGraph checkpoint saver." optional = false python-versions = "^3.9.0,<4.0" @@ -1375,7 +1375,7 @@ files = [] develop = true [package.dependencies] -langgraph-checkpoint = "^2.0.7" +langgraph-checkpoint = "^2.0.10" orjson = ">=3.10.1" psycopg = "^3.2.0" psycopg-pool = "^3.2.0" @@ -1386,7 +1386,7 @@ url = "../checkpoint-postgres" [[package]] name = "langgraph-checkpoint-sqlite" -version = "2.0.2" +version = "2.0.3" description = "Library with a SQLite implementation of LangGraph checkpoint saver." optional = false python-versions = "^3.9.0" @@ -1396,7 +1396,7 @@ develop = true [package.dependencies] aiosqlite = "^0.20.0" -langgraph-checkpoint = "^2.0.2" +langgraph-checkpoint = "^2.0.10" [package.source] type = "directory" @@ -3491,4 +3491,4 @@ type = ["pytest-mypy"] [metadata] lock-version = "2.1" python-versions = ">=3.9.0,<4.0" -content-hash = "356f7e84cf1375119bd3a118ecf4e9476d95913736f8c426a76d55717eb39161" +content-hash = "caf943b02b6913c05d15c37fda6d216669f789e2a059b7e8e2490b2bdcd23e0e" diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index bab202053..04c338c56 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -10,7 +10,7 @@ repository = "https://www.github.com/langchain-ai/langgraph" [tool.poetry.dependencies] python = ">=3.9.0,<4.0" langchain-core = ">=0.2.43,<0.4.0,!=0.3.0,!=0.3.1,!=0.3.2,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14,!=0.3.15,!=0.3.16,!=0.3.17,!=0.3.18,!=0.3.19,!=0.3.20,!=0.3.21,!=0.3.22" -langgraph-checkpoint = "^2.0.4" +langgraph-checkpoint = "^2.0.10" langgraph-sdk = "^0.1.42" [tool.poetry.group.dev.dependencies] diff --git a/libs/langgraph/tests/test_algo.py b/libs/langgraph/tests/test_algo.py index 65a2d28c1..a3a588f6b 100644 --- a/libs/langgraph/tests/test_algo.py +++ b/libs/langgraph/tests/test_algo.py @@ -1,6 +1,6 @@ from langgraph.checkpoint.base import empty_checkpoint from langgraph.constants import PULL, PUSH -from langgraph.pregel.algo import _tuple_str, prepare_next_tasks +from langgraph.pregel.algo import prepare_next_tasks, task_path_str from langgraph.pregel.manager import ChannelsManager @@ -49,16 +49,16 @@ def test_tuple_str() -> None: push_path_b = (PUSH, push_path_a, 1) push_path_c = (PUSH, push_path_b, 3) - assert _tuple_str(push_path_a) == f"~{PUSH}, 0000000002" - assert _tuple_str(push_path_b) == f"~{PUSH}, ~{PUSH}, 0000000002, 0000000001" + assert task_path_str(push_path_a) == f"~{PUSH}, 0000000002" + assert task_path_str(push_path_b) == f"~{PUSH}, ~{PUSH}, 0000000002, 0000000001" assert ( - _tuple_str(push_path_c) + task_path_str(push_path_c) == f"~{PUSH}, ~{PUSH}, ~{PUSH}, 0000000002, 0000000001, 0000000003" ) - assert _tuple_str(pull_path_a) == f"~{PULL}, abc" + assert task_path_str(pull_path_a) == f"~{PULL}, abc" path_list = [push_path_b, push_path_a, pull_path_a, push_path_c] - assert sorted(map(_tuple_str, path_list)) == [ + assert sorted(map(task_path_str, path_list)) == [ f"~{PULL}, abc", f"~{PUSH}, 0000000002", f"~{PUSH}, ~{PUSH}, 0000000002, 0000000001", diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index 15773df47..4d2c45234 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -7255,7 +7255,6 @@ def test_branch_then( ) -@pytest.mark.skip("TODO: re-enable in next PR") @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_send_dedupe_on_resume( request: pytest.FixtureRequest, checkpointer_name: str @@ -7314,30 +7313,35 @@ def test_send_dedupe_on_resume( assert graph.invoke(["0"], thread1, debug=1) == [ "0", "1", + "3.1", "2|Command(goto=Send(node='2', arg=3))", "2|Command(goto=Send(node='flaky', arg=4))", + "3", "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) + if "shallow" in checkpointer_name: + pytest.xfail("TODO: shallow checkpointer reports wrong next set") assert state.next == ("flaky",) # check history if "shallow" not in checkpointer_name: history = [c for c in graph.get_state_history(thread1)] - assert len(history) == 2 + assert len(history) == 4 # resume execution assert graph.invoke(None, thread1, debug=1) == [ "0", "1", + "3.1", "2|Command(goto=Send(node='2', arg=3))", "2|Command(goto=Send(node='flaky', arg=4))", + "3", "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 @@ -7353,12 +7357,13 @@ def test_send_dedupe_on_resume( values=[ "0", "1", + "3.1", "2|Command(goto=Send(node='2', arg=3))", "2|Command(goto=Send(node='flaky', arg=4))", + "3", "2|3", "flaky|4", "3", - "3.1", ], next=(), config={ @@ -7370,35 +7375,33 @@ def test_send_dedupe_on_resume( }, metadata={ "source": "loop", - "writes": {"3": ["3"], "3.1": ["3.1"]}, + "writes": {"3": ["3"]}, "thread_id": "1", - "step": 2, + "step": 4, "parents": {}, }, created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), } - ), + }, tasks=(), ), StateSnapshot( values=[ "0", "1", + "3.1", "2|Command(goto=Send(node='2', arg=3))", "2|Command(goto=Send(node='flaky', arg=4))", + "3", "2|3", "flaky|4", ], - next=("3", "3.1"), + next=("3",), config={ "configurable": { "thread_id": "1", @@ -7408,17 +7411,9 @@ def test_send_dedupe_on_resume( }, metadata={ "source": "loop", - "writes": { - "1": ["1"], - "2": [ - ["2|Command(goto=Send(node='2', arg=3))"], - ["2|Command(goto=Send(node='flaky', arg=4))"], - ["2|3"], - ], - "flaky": ["flaky|4"], - }, + "writes": {"2": ["2|3"], "3": ["3"], "flaky": ["flaky|4"]}, "thread_id": "1", - "step": 1, + "step": 3, "parents": {}, }, created_at=AnyStr(), @@ -7439,6 +7434,123 @@ def test_send_dedupe_on_resume( state=None, result=["3"], ), + ), + ), + StateSnapshot( + values=[ + "0", + "1", + "3.1", + "2|Command(goto=Send(node='2', arg=3))", + "2|Command(goto=Send(node='flaky', arg=4))", + ], + next=("2", "flaky", "3"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "2": [ + ["2|Command(goto=Send(node='2', arg=3))"], + ["2|Command(goto=Send(node='flaky', arg=4))"], + ], + "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=( + PregelTask( + id=AnyStr(), + name="2", + path=("__pregel_push", 0), + error=None, + interrupts=(), + state=None, + result=["2|3"], + ), + PregelTask( + id=AnyStr(), + name="flaky", + path=("__pregel_push", 1), + error=None, + interrupts=( + Interrupt( + value="Bahh", resumable=False, ns=None, when="during" + ), + ), + state=None, + result=["flaky|4"], + ), + PregelTask( + id=AnyStr(), + name="3", + path=("__pregel_pull", "3"), + error=None, + interrupts=(), + state=None, + result=["3"], + ), + ), + ), + StateSnapshot( + values=["0", "1"], + next=("2", "2", "3.1"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"1": ["1"]}, + "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="2", + path=("__pregel_push", 0), + error=None, + interrupts=(), + state=None, + result=["2|Command(goto=Send(node='2', arg=3))"], + ), + PregelTask( + id=AnyStr(), + name="2", + path=("__pregel_push", 1), + error=None, + interrupts=(), + state=None, + result=["2|Command(goto=Send(node='flaky', arg=4))"], + ), PregelTask( id=AnyStr(), name="3.1", @@ -7452,7 +7564,7 @@ def test_send_dedupe_on_resume( ), StateSnapshot( values=["0"], - next=("1", "2", "2", "2", "flaky"), + next=("1",), config={ "configurable": { "thread_id": "1", @@ -7485,66 +7597,6 @@ def test_send_dedupe_on_resume( state=None, result=["1"], ), - PregelTask( - id=AnyStr(), - name="2", - path=( - "__pregel_push", - ("__pregel_pull", "1"), - 2, - ), - error=None, - interrupts=(), - state=None, - result=["2|Command(goto=Send(node='2', arg=3))"], - ), - PregelTask( - id=AnyStr(), - name="2", - path=( - "__pregel_push", - ("__pregel_pull", "1"), - 3, - ), - error=None, - interrupts=(), - state=None, - result=["2|Command(goto=Send(node='flaky', arg=4))"], - ), - PregelTask( - id=AnyStr(), - name="2", - path=( - "__pregel_push", - ( - "__pregel_push", - ("__pregel_pull", "1"), - 2, - ), - 2, - ), - error=None, - interrupts=(), - state=None, - result=["2|3"], - ), - PregelTask( - id=AnyStr(), - name="flaky", - path=( - "__pregel_push", - ( - "__pregel_push", - ("__pregel_pull", "1"), - 3, - ), - 2, - ), - error=None, - interrupts=(Interrupt(value="Bahh", when="during"),), - state=None, - result=["flaky|4"], - ), ), ), StateSnapshot( diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 427ee4041..c05682f7e 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -2373,7 +2373,6 @@ async def test_concurrent_emit_sends() -> None: ) -@pytest.mark.skip("TODO: re-enable in next PR") @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_send_sequences(checkpointer_name: str) -> None: class Node: @@ -2621,7 +2620,6 @@ async def test_imp_stream_order(checkpointer_name: str) -> None: ] -@pytest.mark.skip("TODO: re-enable in next PR") @pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: class InterruptOnce: @@ -2685,7 +2683,6 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: ] assert builder.nodes["2"].runnable.func.ticks == 3 assert builder.nodes["flaky"].runnable.func.ticks == 1 - print((await graph.aget_state(thread1)).tasks) # resume execution assert await graph.ainvoke(None, thread1, debug=1) == [ "0", @@ -2694,8 +2691,8 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: "2|Command(goto=Send(node='2', arg=3))", "2|Command(goto=Send(node='flaky', arg=4))", "3", - "flaky|4", "2|3", + "flaky|4", "3", ] # node "2" doesn't get called again, as we recover writes saved before @@ -2713,8 +2710,8 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: "2|Command(goto=Send(node='2', arg=3))", "2|Command(goto=Send(node='flaky', arg=4))", "3", - "flaky|4", "2|3", + "flaky|4", "3", ], next=(), @@ -2750,8 +2747,8 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: "2|Command(goto=Send(node='2', arg=3))", "2|Command(goto=Send(node='flaky', arg=4))", "3", - "flaky|4", "2|3", + "flaky|4", ], next=("3",), config={