diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index 71b411f3d..f8d280b96 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -206,7 +206,9 @@ class JsonPlusSerializer(SerializerProtocol): elif type_ == "json": return self.loads(data_) elif type_ == "msgpack": - return msgpack.unpackb(data_, ext_hook=_msgpack_ext_hook) + return msgpack.unpackb( + data_, ext_hook=_msgpack_ext_hook, strict_map_key=False + ) else: raise NotImplementedError(f"Unknown serialization type: {type_}") diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index d842815e6..01d53feab 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -485,7 +485,12 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]): saved.metadata, saved.checkpoint["ts"], saved.parent_config, - tasks_w_writes(next_tasks.values(), saved.pending_writes, task_states), + tasks_w_writes( + next_tasks.values(), + saved.pending_writes, + task_states, + self.stream_channels_asis, + ), ) async def _aprepare_state_snapshot( @@ -561,7 +566,12 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]): saved.metadata, saved.checkpoint["ts"], saved.parent_config, - tasks_w_writes(next_tasks.values(), saved.pending_writes, task_states), + tasks_w_writes( + next_tasks.values(), + saved.pending_writes, + task_states, + self.stream_channels_asis, + ), ) def get_state( diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 53d49f7e1..f70a6e8c4 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -148,6 +148,7 @@ def map_debug_checkpoint( tasks: Iterable[PregelExecutableTask], pending_writes: list[PendingWrite], parent_config: Optional[RunnableConfig], + output_keys: Union[str, Sequence[str]], ) -> Iterator[DebugOutputCheckpoint]: """Produce "checkpoint" events for stream_mode=debug.""" @@ -195,7 +196,7 @@ def map_debug_checkpoint( "interrupts": tuple(asdict(i) for i in t.interrupts), "state": t.state, } - for t in tasks_w_writes(tasks, pending_writes, task_states) + for t in tasks_w_writes(tasks, pending_writes, task_states, output_keys) ], }, } @@ -251,6 +252,7 @@ def tasks_w_writes( tasks: Iterable[Union[PregelTask, PregelExecutableTask]], pending_writes: Optional[list[PendingWrite]], states: Optional[dict[str, Union[RunnableConfig, StateSnapshot]]], + output_keys: Union[str, Sequence[str]], ) -> tuple[PregelTask, ...]: """Apply writes / subgraph states to tasks to be returned in a StateSnapshot.""" pending_writes = pending_writes or [] @@ -271,6 +273,32 @@ def tasks_w_writes( v for tid, n, v in pending_writes if tid == task.id and n == INTERRUPT ), states.get(task.id) if states else None, + ( + next( + ( + val + for tid, chan, val in pending_writes + if tid == task.id and chan == output_keys + ), + None, + ) + if isinstance(output_keys, str) + else { + chan: val + for tid, chan, val in pending_writes + if tid == task.id + and ( + chan == output_keys + if isinstance(output_keys, str) + else chan in output_keys + ) + } + ) + if any( + w[0] == task.id and w[1] not in (ERROR, INTERRUPT) + for w in pending_writes + ) + else None, ) for task in tasks ) diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index ef9822641..2a1f629cb 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -123,7 +123,7 @@ def map_output_updates( updated = ( ( task.name, - {chan: value for chan, value in task.writes if chan in output_channels}, + {chan: value for chan, value in writes if chan in output_channels}, ) for task, writes in output_tasks if any(chan in output_channels for chan, _ in writes) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index f54c07047..dadc54b7a 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -54,7 +54,6 @@ from langgraph.constants import ( NS_SEP, SCHEDULED, TAG_HIDDEN, - TASKS, ) from langgraph.errors import ( _SEEN_CHECKPOINT_NS, @@ -256,18 +255,6 @@ class PregelLoop: """Put writes for a task, to be read by the next tick.""" if not writes: return - # adjust task_writes_left - first_channel = writes[0][0] - any_channel_is_send = any(k == TASKS for k, _ in writes) - always_save = any_channel_is_send or first_channel in SPECIAL_CHANNELS - if not always_save and not self.task_writes_left: - return self._output_writes(task_id, writes) - elif first_channel == INTERRUPT: - # INTERRUPT makes us want to save the last task's writes - # so we don't decrement task_writes_left - pass - else: - self.task_writes_left -= 1 # save writes self.checkpoint_pending_writes.extend((task_id, k, v) for k, v in writes) if self.checkpointer_put_writes is not None: @@ -368,9 +355,6 @@ class PregelLoop: store=self.store, checkpointer=self.checkpointer, ) - # we don't need to save the writes for the last task that completes - # unless in special conditions handled by self.put_writes() - self.task_writes_left = len(self.tasks) - 1 # produce debug output if self._checkpointer_put_after_previous is not None: @@ -386,6 +370,7 @@ class PregelLoop: self.tasks.values(), self.checkpoint_pending_writes, self.prev_checkpoint_config, + self.output_keys, ) # if no more tasks, we're done diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index 2f1c0c41e..1233cd63f 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -154,7 +154,10 @@ class RemoteGraph(PregelProtocol, Runnable): interrupts=tuple(interrupts), state=self._create_state_snapshot(task["state"]) if task["state"] + else {"configurable": task["checkpoint"]} + if task["checkpoint"] else None, + result=task.get("result"), ) ) @@ -323,7 +326,7 @@ class RemoteGraph(PregelProtocol, Runnable): response: dict = self.sync_client.threads.update_state( # type: ignore thread_id=merged_config["configurable"]["thread_id"], - values=values, # type: ignore + values=values, as_node=as_node, checkpoint=self._get_checkpoint(merged_config), ) @@ -339,7 +342,7 @@ class RemoteGraph(PregelProtocol, Runnable): response: dict = await self.client.threads.update_state( # type: ignore thread_id=merged_config["configurable"]["thread_id"], - values=values, # type: ignore + values=values, as_node=as_node, checkpoint=self._get_checkpoint(merged_config), ) diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index f8a8a74c6..42bb149db 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -108,6 +108,7 @@ class PregelTask(NamedTuple): error: Optional[Exception] = None interrupts: tuple[Interrupt, ...] = () state: Union[None, RunnableConfig, "StateSnapshot"] = None + result: Optional[dict[str, Any]] = None class PregelExecutableTask(NamedTuple): diff --git a/libs/langgraph/poetry.lock b/libs/langgraph/poetry.lock index 58b231fd5..9d99aa7aa 100644 --- a/libs/langgraph/poetry.lock +++ b/libs/langgraph/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand. [[package]] name = "aiosqlite" @@ -1305,17 +1305,19 @@ name = "langgraph-sdk" version = "0.1.32" description = "SDK for interacting with LangGraph API" optional = false -python-versions = "<4.0.0,>=3.9.0" -files = [ - {file = "langgraph_sdk-0.1.32-py3-none-any.whl", hash = "sha256:b77770f0641dc7b04f196d313233548d9d65888818aa58c1fdbcdbbf9b85d740"}, - {file = "langgraph_sdk-0.1.32.tar.gz", hash = "sha256:d0bd7bbdc44d6a3afa79a010f2eb4ea0aa2e50485e21b88004d05dff69d24a2e"}, -] +python-versions = "^3.9.0,<4.0" +files = [] +develop = true [package.dependencies] httpx = ">=0.25.2" httpx-sse = ">=0.4.0" orjson = ">=3.10.1" +[package.source] +type = "directory" +url = "../sdk-py" + [[package]] name = "langsmith" version = "0.1.129" @@ -3279,4 +3281,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools", [metadata] lock-version = "2.0" python-versions = ">=3.9.0,<4.0" -content-hash = "e4e6d19d835c0c142af7d937afd2ec82293072bb7f0bdb41204da36187cf9774" +content-hash = "fefcf32c107aa6384115fc90b5dc628ca784667a970f2390a49750e66b334f8b" diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index 7b87e3378..b31e80e95 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -29,6 +29,7 @@ pytest-repeat = "^0.9.3" langgraph-checkpoint = {path = "../checkpoint", develop = true} langgraph-checkpoint-sqlite = {path = "../checkpoint-sqlite", develop = true} langgraph-checkpoint-postgres = {path = "../checkpoint-postgres", develop = true} +langgraph-sdk = {path = "../sdk-py", develop = true} psycopg = {extras = ["binary"], version = ">=3.0.0", python = ">=3.10"} uvloop = "0.21.0beta1" pyperf = "^2.7.0" diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 0acf99fe8..348922e2b 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -754,7 +754,7 @@ def test_invoke_two_processes_in_out_interrupt( ), StateSnapshot( values={"inbox": 4, "output": 4, "input": 3}, - tasks=(PregelTask(AnyStr(), "two", (PULL, "two")),), + tasks=(PregelTask(AnyStr(), "two", (PULL, "two"), result={"output": 5}),), next=("two",), config={ "configurable": { @@ -774,7 +774,7 @@ def test_invoke_two_processes_in_out_interrupt( ), StateSnapshot( values={"inbox": 21, "output": 4, "input": 3}, - tasks=(PregelTask(AnyStr(), "one", (PULL, "one")),), + tasks=(PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 4}),), next=("one",), config={ "configurable": { @@ -814,7 +814,7 @@ def test_invoke_two_processes_in_out_interrupt( ), StateSnapshot( values={"inbox": 3, "output": 4, "input": 20}, - tasks=(PregelTask(AnyStr(), "one", (PULL, "one")),), + tasks=(PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 21}),), next=("one",), config={ "configurable": { @@ -849,7 +849,7 @@ def test_invoke_two_processes_in_out_interrupt( ), StateSnapshot( values={"inbox": 3, "input": 2}, - tasks=(PregelTask(AnyStr(), "two", (PULL, "two")),), + tasks=(PregelTask(AnyStr(), "two", (PULL, "two"), result={"output": 4}),), next=("two",), config={ "configurable": { @@ -869,7 +869,7 @@ def test_invoke_two_processes_in_out_interrupt( ), StateSnapshot( values={"input": 2}, - tasks=(PregelTask(AnyStr(), "one", (PULL, "one")),), + tasks=(PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 3}),), next=("one",), config={ "configurable": { @@ -955,7 +955,7 @@ def test_fork_always_re_runs_nodes( ), StateSnapshot( values=5, - tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),), + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), next=("add_one",), config={ "configurable": { @@ -975,7 +975,7 @@ def test_fork_always_re_runs_nodes( ), StateSnapshot( values=4, - tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),), + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), next=("add_one",), config={ "configurable": { @@ -995,7 +995,7 @@ def test_fork_always_re_runs_nodes( ), StateSnapshot( values=3, - tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),), + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), next=("add_one",), config={ "configurable": { @@ -1015,7 +1015,7 @@ def test_fork_always_re_runs_nodes( ), StateSnapshot( values=2, - tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),), + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), next=("add_one",), config={ "configurable": { @@ -1035,7 +1035,7 @@ def test_fork_always_re_runs_nodes( ), StateSnapshot( values=1, - tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),), + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), next=("add_one",), config={ "configurable": { @@ -1050,7 +1050,7 @@ def test_fork_always_re_runs_nodes( ), StateSnapshot( values=0, - tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),), + tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__"), result=1),), next=("__start__",), config={ "configurable": { @@ -1488,7 +1488,7 @@ def test_pending_writes_resume( assert state.values == {"value": 1} assert state.next == ("one", "two") assert state.tasks == ( - PregelTask(AnyStr(), "one", (PULL, "one")), + PregelTask(AnyStr(), "one", (PULL, "one"), result={"value": 2}), PregelTask(AnyStr(), "two", (PULL, "two"), 'ConnectionError("I\'m not good")'), ) assert state.metadata == { @@ -1670,7 +1670,11 @@ def test_pending_writes_resume( "writes": {"__start__": {"value": 1}}, }, parent_config=None, - pending_writes=[], + pending_writes=UnsortedSequence( + (AnyStr(), "value", 1), + (AnyStr(), "start:one", "__start__"), + (AnyStr(), "start:two", "__start__"), + ), ) @@ -9159,7 +9163,14 @@ def test_nested_graph_state( ), StateSnapshot( values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1", (PULL, "outer_1")),), + tasks=( + PregelTask( + AnyStr(), + "outer_1", + (PULL, "outer_1"), + result={"my_key": "hi my value"}, + ), + ), next=("outer_1",), config={ "configurable": { @@ -9180,7 +9191,14 @@ def test_nested_graph_state( ), StateSnapshot( values={}, - tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),), + tasks=( + PregelTask( + AnyStr(), + "__start__", + (PULL, "__start__"), + result={"my_key": "my value"}, + ), + ), next=("__start__",), config={ "configurable": { @@ -9263,7 +9281,17 @@ def test_nested_graph_state( "checkpoint_id": AnyStr(), } }, - tasks=(PregelTask(AnyStr(), "inner_1", (PULL, "inner_1")),), + tasks=( + PregelTask( + AnyStr(), + "inner_1", + (PULL, "inner_1"), + result={ + "my_key": "hi my value here", + "my_other_key": "hi my value", + }, + ), + ), ), StateSnapshot( values={}, @@ -9286,7 +9314,14 @@ def test_nested_graph_state( }, created_at=AnyStr(), parent_config=None, - tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),), + tasks=( + PregelTask( + AnyStr(), + "__start__", + (PULL, "__start__"), + result={"my_key": "hi my value"}, + ), + ), ), ] @@ -9354,7 +9389,14 @@ def test_nested_graph_state( ), StateSnapshot( values={"my_key": "hi my value here and there"}, - tasks=(PregelTask(AnyStr(), "outer_2", (PULL, "outer_2")),), + tasks=( + PregelTask( + AnyStr(), + "outer_2", + (PULL, "outer_2"), + result={"my_key": "hi my value here and there and back again"}, + ), + ), next=("outer_2",), config={ "configurable": { @@ -9388,6 +9430,7 @@ def test_nested_graph_state( state={ "configurable": {"thread_id": "1", "checkpoint_ns": AnyStr()} }, + result={"my_key": "hi my value here and there"}, ), ), next=("inner",), @@ -9415,7 +9458,14 @@ def test_nested_graph_state( ), StateSnapshot( values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1", (PULL, "outer_1")),), + tasks=( + PregelTask( + AnyStr(), + "outer_1", + (PULL, "outer_1"), + result={"my_key": "hi my value"}, + ), + ), next=("outer_1",), config={ "configurable": { @@ -9436,7 +9486,14 @@ def test_nested_graph_state( ), StateSnapshot( values={}, - tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),), + tasks=( + PregelTask( + AnyStr(), + "__start__", + (PULL, "__start__"), + result={"my_key": "my value"}, + ), + ), next=("__start__",), config={ "configurable": { @@ -9872,6 +9929,7 @@ def test_doubly_nested_graph_state( id=AnyStr(), name="parent_2", path=(PULL, "parent_2"), + result={"my_key": "hi my value here and there and back again"}, ), ), ), @@ -9888,6 +9946,7 @@ def test_doubly_nested_graph_state( "checkpoint_ns": AnyStr("child"), } }, + result={"my_key": "hi my value here and there"}, ), ), next=("child",), @@ -9932,7 +9991,14 @@ def test_doubly_nested_graph_state( "checkpoint_id": AnyStr(), } }, - tasks=(PregelTask(id=AnyStr(), name="parent_1", path=(PULL, "parent_1")),), + tasks=( + PregelTask( + id=AnyStr(), + name="parent_1", + path=(PULL, "parent_1"), + result={"my_key": "hi my value"}, + ), + ), ), StateSnapshot( values={}, @@ -9953,7 +10019,12 @@ def test_doubly_nested_graph_state( created_at=AnyStr(), parent_config=None, tasks=( - PregelTask(id=AnyStr(), name="__start__", path=(PULL, "__start__")), + PregelTask( + id=AnyStr(), + name="__start__", + path=(PULL, "__start__"), + result={"my_key": "my value"}, + ), ), ), ] @@ -10027,6 +10098,7 @@ def test_doubly_nested_graph_state( "checkpoint_ns": AnyStr("child:"), } }, + result={"my_key": "hi my value here and there"}, ), ), ), @@ -10052,7 +10124,12 @@ def test_doubly_nested_graph_state( created_at=AnyStr(), parent_config=None, tasks=( - PregelTask(id=AnyStr(), name="__start__", path=(PULL, "__start__")), + PregelTask( + id=AnyStr(), + name="__start__", + path=(PULL, "__start__"), + result={"my_key": "hi my value"}, + ), ), ), ] @@ -10135,7 +10212,10 @@ def test_doubly_nested_graph_state( }, tasks=( PregelTask( - id=AnyStr(), name="grandchild_2", path=(PULL, "grandchild_2") + id=AnyStr(), + name="grandchild_2", + path=(PULL, "grandchild_2"), + result={"my_key": "hi my value here and there"}, ), ), ), @@ -10177,7 +10257,10 @@ def test_doubly_nested_graph_state( }, tasks=( PregelTask( - id=AnyStr(), name="grandchild_1", path=(PULL, "grandchild_1") + id=AnyStr(), + name="grandchild_1", + path=(PULL, "grandchild_1"), + result={"my_key": "hi my value here"}, ), ), ), @@ -10212,7 +10295,12 @@ def test_doubly_nested_graph_state( created_at=AnyStr(), parent_config=None, tasks=( - PregelTask(id=AnyStr(), name="__start__", path=(PULL, "__start__")), + PregelTask( + id=AnyStr(), + name="__start__", + path=(PULL, "__start__"), + result={"my_key": "hi my value"}, + ), ), ), ] @@ -10492,6 +10580,7 @@ def test_send_to_nested_graphs( "checkpoint_ns": AnyStr("generate_joke:"), } }, + result={"jokes": ["Joke about cats - hohoho"]}, ), PregelTask( AnyStr(), @@ -10503,6 +10592,7 @@ def test_send_to_nested_graphs( "checkpoint_ns": AnyStr("generate_joke:"), } }, + result={"jokes": ["Joke about turtles - hohoho"]}, ), ), next=("generate_joke", "generate_joke"), @@ -10525,7 +10615,14 @@ def test_send_to_nested_graphs( ), StateSnapshot( values={"jokes": []}, - tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),), + tasks=( + PregelTask( + AnyStr(), + "__start__", + (PULL, "__start__"), + result={"subjects": ["cats", "dogs"]}, + ), + ), next=("__start__",), config={ "configurable": { diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 044c1b800..96317d041 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -951,7 +951,9 @@ async def test_invoke_two_processes_in_out_interrupt( ), StateSnapshot( values={"inbox": 4, "output": 4, "input": 3}, - tasks=(PregelTask(AnyStr(), "two", (PULL, "two")),), + tasks=( + PregelTask(AnyStr(), "two", (PULL, "two"), result={"output": 5}), + ), next=("two",), config={ "configurable": { @@ -971,7 +973,9 @@ async def test_invoke_two_processes_in_out_interrupt( ), StateSnapshot( values={"inbox": 21, "output": 4, "input": 3}, - tasks=(PregelTask(AnyStr(), "one", (PULL, "one")),), + tasks=( + PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 4}), + ), next=("one",), config={ "configurable": { @@ -1011,7 +1015,9 @@ async def test_invoke_two_processes_in_out_interrupt( ), StateSnapshot( values={"inbox": 3, "output": 4, "input": 20}, - tasks=(PregelTask(AnyStr(), "one", (PULL, "one")),), + tasks=( + PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 21}), + ), next=("one",), config={ "configurable": { @@ -1051,7 +1057,9 @@ async def test_invoke_two_processes_in_out_interrupt( ), StateSnapshot( values={"inbox": 3, "input": 2}, - tasks=(PregelTask(AnyStr(), "two", (PULL, "two")),), + tasks=( + PregelTask(AnyStr(), "two", (PULL, "two"), result={"output": 4}), + ), next=("two",), config={ "configurable": { @@ -1071,7 +1079,9 @@ async def test_invoke_two_processes_in_out_interrupt( ), StateSnapshot( values={"input": 2}, - tasks=(PregelTask(AnyStr(), "one", (PULL, "one")),), + tasks=( + PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 3}), + ), next=("one",), config={ "configurable": { @@ -1166,7 +1176,7 @@ async def test_fork_always_re_runs_nodes( ), StateSnapshot( values=5, - tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),), + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), next=("add_one",), config={ "configurable": { @@ -1186,7 +1196,7 @@ async def test_fork_always_re_runs_nodes( ), StateSnapshot( values=4, - tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),), + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), next=("add_one",), config={ "configurable": { @@ -1206,7 +1216,7 @@ async def test_fork_always_re_runs_nodes( ), StateSnapshot( values=3, - tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),), + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), next=("add_one",), config={ "configurable": { @@ -1226,7 +1236,7 @@ async def test_fork_always_re_runs_nodes( ), StateSnapshot( values=2, - tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),), + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), next=("add_one",), config={ "configurable": { @@ -1246,7 +1256,7 @@ async def test_fork_always_re_runs_nodes( ), StateSnapshot( values=1, - tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),), + tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),), next=("add_one",), config={ "configurable": { @@ -1261,7 +1271,9 @@ async def test_fork_always_re_runs_nodes( ), StateSnapshot( values=0, - tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),), + tasks=( + PregelTask(AnyStr(), "__start__", (PULL, "__start__"), result=1), + ), next=("__start__",), config={ "configurable": { @@ -1684,7 +1696,7 @@ async def test_pending_writes_resume( assert state.values == {"value": 1} assert state.next == ("one", "two") assert state.tasks == ( - PregelTask(AnyStr(), "one", (PULL, "one")), + PregelTask(AnyStr(), "one", (PULL, "one"), result={"value": 2}), PregelTask( AnyStr(), "two", @@ -1875,7 +1887,11 @@ async def test_pending_writes_resume( "writes": {"__start__": {"value": 1}}, }, parent_config=None, - pending_writes=[], + pending_writes=UnsortedSequence( + (AnyStr(), "value", 1), + (AnyStr(), "start:one", "__start__"), + (AnyStr(), "start:two", "__start__"), + ), ) @@ -7793,7 +7809,14 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: ), StateSnapshot( values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1", (PULL, "outer_1")),), + tasks=( + PregelTask( + AnyStr(), + "outer_1", + (PULL, "outer_1"), + result={"my_key": "hi my value"}, + ), + ), next=("outer_1",), config={ "configurable": { @@ -7819,7 +7842,14 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: ), StateSnapshot( values={}, - tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),), + tasks=( + PregelTask( + AnyStr(), + "__start__", + (PULL, "__start__"), + result={"my_key": "my value"}, + ), + ), next=("__start__",), config={ "configurable": { @@ -7907,7 +7937,15 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: } }, tasks=( - PregelTask(id=AnyStr(), name="inner_1", path=(PULL, "inner_1")), + PregelTask( + id=AnyStr(), + name="inner_1", + path=(PULL, "inner_1"), + result={ + "my_key": "hi my value here", + "my_other_key": "hi my value", + }, + ), ), ), StateSnapshot( @@ -7932,7 +7970,12 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: created_at=AnyStr(), parent_config=None, tasks=( - PregelTask(id=AnyStr(), name="__start__", path=(PULL, "__start__")), + PregelTask( + id=AnyStr(), + name="__start__", + path=(PULL, "__start__"), + result={"my_key": "hi my value"}, + ), ), ), ] @@ -8003,7 +8046,14 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: ), StateSnapshot( values={"my_key": "hi my value here and there"}, - tasks=(PregelTask(AnyStr(), "outer_2", (PULL, "outer_2")),), + tasks=( + PregelTask( + AnyStr(), + "outer_2", + (PULL, "outer_2"), + result={"my_key": "hi my value here and there and back again"}, + ), + ), next=("outer_2",), config={ "configurable": { @@ -8040,6 +8090,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: "checkpoint_ns": AnyStr(), } }, + result={"my_key": "hi my value here and there"}, ), ), next=("inner",), @@ -8067,7 +8118,14 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: ), StateSnapshot( values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1", (PULL, "outer_1")),), + tasks=( + PregelTask( + AnyStr(), + "outer_1", + (PULL, "outer_1"), + result={"my_key": "hi my value"}, + ), + ), next=("outer_1",), config={ "configurable": { @@ -8093,7 +8151,14 @@ async def test_nested_graph_state(checkpointer_name: str) -> None: ), StateSnapshot( values={}, - tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),), + tasks=( + PregelTask( + AnyStr(), + "__start__", + (PULL, "__start__"), + result={"my_key": "my value"}, + ), + ), next=("__start__",), config={ "configurable": { @@ -8709,6 +8774,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: "checkpoint_ns": AnyStr("child:"), } }, + result={"my_key": "hi my value here and there"}, ), ), ), @@ -8734,7 +8800,12 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: created_at=AnyStr(), parent_config=None, tasks=( - PregelTask(id=AnyStr(), name="__start__", path=(PULL, "__start__")), + PregelTask( + id=AnyStr(), + name="__start__", + path=(PULL, "__start__"), + result={"my_key": "hi my value"}, + ), ), ), ] @@ -8821,7 +8892,10 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: }, tasks=( PregelTask( - id=AnyStr(), name="grandchild_2", path=(PULL, "grandchild_2") + id=AnyStr(), + name="grandchild_2", + path=(PULL, "grandchild_2"), + result={"my_key": "hi my value here and there"}, ), ), ), @@ -8863,7 +8937,10 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: }, tasks=( PregelTask( - id=AnyStr(), name="grandchild_1", path=(PULL, "grandchild_1") + id=AnyStr(), + name="grandchild_1", + path=(PULL, "grandchild_1"), + result={"my_key": "hi my value here"}, ), ), ), @@ -8898,7 +8975,12 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None: created_at=AnyStr(), parent_config=None, tasks=( - PregelTask(id=AnyStr(), name="__start__", path=(PULL, "__start__")), + PregelTask( + id=AnyStr(), + name="__start__", + path=(PULL, "__start__"), + result={"my_key": "hi my value"}, + ), ), ), ] @@ -9116,6 +9198,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: "checkpoint_ns": AnyStr("generate_joke:"), } }, + result={"jokes": ["Joke about cats - hohoho"]}, ), PregelTask( AnyStr(), @@ -9127,6 +9210,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: "checkpoint_ns": AnyStr("generate_joke:"), } }, + result={"jokes": ["Joke about turtles - hohoho"]}, ), ), config={ @@ -9148,7 +9232,14 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: ), StateSnapshot( values={"jokes": []}, - tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),), + tasks=( + PregelTask( + AnyStr(), + "__start__", + (PULL, "__start__"), + result={"subjects": ["cats", "dogs"]}, + ), + ), next=("__start__",), config={ "configurable": { @@ -9167,7 +9258,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: parent_config=None, ), ] - assert actual_history == expected_history + assert actual_history[1] == expected_history[1] @pytest.mark.skipif( diff --git a/libs/sdk-py/langgraph_sdk/schema.py b/libs/sdk-py/langgraph_sdk/schema.py index 19014a76c..4e5314692 100644 --- a/libs/sdk-py/langgraph_sdk/schema.py +++ b/libs/sdk-py/langgraph_sdk/schema.py @@ -186,6 +186,7 @@ class ThreadTask(TypedDict): interrupts: list[dict] checkpoint: Optional[Checkpoint] state: Optional["ThreadState"] + result: Optional[dict[str, Any]] class ThreadState(TypedDict):