diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index e2d9f069a..b9d6f37f7 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -46,6 +46,8 @@ CONFIG_KEY_SEND = sys.intern("__pregel_send") # holds the `write` function that accepts writes to state/edges/reserved keys CONFIG_KEY_READ = sys.intern("__pregel_read") # holds the `read` function that returns a copy of the current state +CONFIG_KEY_CALL = sys.intern("__pregel_call") +# holds the `call` function that accepts a node/func, args and returns a future CONFIG_KEY_CHECKPOINTER = sys.intern("__pregel_checkpointer") # holds a `BaseCheckpointSaver` passed from parent graph to child graphs CONFIG_KEY_STREAM = sys.intern("__pregel_stream") diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index e714afe21..52cd2bf09 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -541,6 +541,11 @@ class Pregel(PregelProtocol): next_tasks[tid].writes.append((k, v)) if tasks := [t for t in next_tasks.values() if t.writes]: apply_writes(saved.checkpoint, channels, tasks, None) + print( + "next_tasks", + [(t.id, t.name, t.input, bool(t.writes)) for t in next_tasks.values()], + ) + print("pending_writes", saved.pending_writes) # assemble the state snapshot return StateSnapshot( read_channels(channels, self.stream_channels_asis), diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 0885f12aa..a8d39674c 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -97,6 +97,11 @@ class PregelTaskWrites(NamedTuple): triggers: Sequence[str] +class Call: + func: str | Callable + input: Any + + def should_interrupt( checkpoint: Checkpoint, interrupt_nodes: Union[All, Sequence[str]], @@ -490,10 +495,16 @@ def prepare_single_task( PUSH, str(idx), ) - elif len(task_path) == 4: + elif len(task_path) >= 4: # new PUSH tasks, executed in superstep n # (PUSH, parent task path, idx of PUSH write, id of parent task) - task_path_t = cast(tuple[str, tuple, int, str], task_path) + task_path_t = cast( + Union[ + tuple[str, tuple, int, str], + tuple[str, tuple, int, str, Optional[Call]], + ], + task_path, + ) writes_for_path = [w for w in pending_writes if w[0] == task_path_t[3]] if task_path_t[2] >= len(writes_for_path): logger.warning( @@ -501,6 +512,13 @@ def prepare_single_task( ) return packet = writes_for_path[task_path_t[2]][2] + if packet is None: + if len(task_path_t) == 5: + packet = task_path_t[4] + else: + # no packet to replay, this is a "call" task + return + # TODO handle Call packets if not isinstance(packet, Send): logger.warning( f"Ignoring invalid packet type {type(packet)} in pending writes" @@ -533,7 +551,7 @@ def prepare_single_task( "langgraph_step": step, "langgraph_node": packet.node, "langgraph_triggers": triggers, - "langgraph_path": task_path, + "langgraph_path": task_path[:3], "langgraph_checkpoint_ns": task_checkpoint_ns, } if task_id_checksum is not None: @@ -572,7 +590,7 @@ def prepare_single_task( channels, managed, PregelTaskWrites( - task_path, packet.node, writes, triggers + task_path[:3], packet.node, writes, triggers ), config, ), @@ -602,12 +620,12 @@ def prepare_single_task( proc.retry_policy, None, task_id, - task_path, + task_path[:3], writers=proc.flat_writers, ) else: - return PregelTask(task_id, packet.node, task_path) + return PregelTask(task_id, packet.node, task_path[:3]) elif task_path[0] == PULL: # (PULL, node name) name = cast(str, task_path[1]) @@ -657,7 +675,7 @@ def prepare_single_task( "langgraph_step": step, "langgraph_node": name, "langgraph_triggers": triggers, - "langgraph_path": task_path, + "langgraph_path": task_path[:3], "langgraph_checkpoint_ns": task_checkpoint_ns, } if task_id_checksum is not None: @@ -696,7 +714,9 @@ def prepare_single_task( checkpoint, channels, managed, - PregelTaskWrites(task_path, name, writes, triggers), + PregelTaskWrites( + task_path[:3], name, writes, triggers + ), config, ), CONFIG_KEY_STORE: ( @@ -725,11 +745,11 @@ def prepare_single_task( proc.retry_policy, None, task_id, - task_path, + task_path[:3], writers=proc.flat_writers, ) else: - return PregelTask(task_id, name, task_path) + return PregelTask(task_id, name, task_path[:3]) def _proc_input( diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index d9af9279e..31397d22b 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -73,6 +73,7 @@ from langgraph.managed.base import ( WritableManagedValue, ) from langgraph.pregel.algo import ( + Call, GetNextVersion, PregelTaskWrites, apply_writes, @@ -307,12 +308,9 @@ class PregelLoop(LoopProtocol): self._output_writes(task_id, writes) def accept_push( - self, task: PregelExecutableTask, write_idx: int + self, task: PregelExecutableTask, write_idx: int, call: Optional[Call] = None ) -> Optional[PregelExecutableTask]: """Accept a PUSH from a task, potentially returning a new task to start.""" - # don't start if an earlier PUSH has already triggered an interrupt - if self.to_interrupt: - return # don't start if we should interrupt *after* the original task if should_interrupt(self.checkpoint, self.interrupt_after, [task]): self.to_interrupt.append(task) @@ -320,7 +318,7 @@ class PregelLoop(LoopProtocol): if pushed := cast( Optional[PregelExecutableTask], prepare_single_task( - (PUSH, task.path, write_idx, task.id), + (PUSH, task.path, write_idx, task.id, call), None, checkpoint=self.checkpoint, pending_writes=[(task.id, *w) for w in task.writes], diff --git a/libs/langgraph/langgraph/pregel/retry.py b/libs/langgraph/langgraph/pregel/retry.py index 2d0f2b6da..093859579 100644 --- a/libs/langgraph/langgraph/pregel/retry.py +++ b/libs/langgraph/langgraph/pregel/retry.py @@ -4,14 +4,12 @@ import random import sys import time from dataclasses import replace -from functools import partial -from typing import Any, Callable, Optional, Sequence +from typing import Any, Optional, Sequence from langgraph.constants import ( CONF, CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_RESUMING, - CONFIG_KEY_SEND, NS_SEP, ) from langgraph.errors import _SEEN_CHECKPOINT_NS, GraphBubbleUp, ParentCommand @@ -25,17 +23,15 @@ SUPPORTS_EXC_NOTES = sys.version_info >= (3, 11) def run_with_retry( task: PregelExecutableTask, retry_policy: Optional[RetryPolicy], - writer: Optional[ - Callable[[PregelExecutableTask, Sequence[tuple[str, Any]]], None] - ] = None, + configurable: Optional[dict[str, Any]] = None, ) -> None: """Run a task with retries.""" retry_policy = task.retry_policy or retry_policy interval = retry_policy.initial_interval if retry_policy else 0 attempts = 0 config = task.config - if writer is not None: - config = patch_configurable(config, {CONFIG_KEY_SEND: partial(writer, task)}) + if configurable is not None: + config = patch_configurable(config, configurable) while True: try: # clear any writes from previous attempts @@ -115,17 +111,15 @@ async def arun_with_retry( task: PregelExecutableTask, retry_policy: Optional[RetryPolicy], stream: bool = False, - writer: Optional[ - Callable[[PregelExecutableTask, Sequence[tuple[str, Any]]], None] - ] = None, + configurable: Optional[dict[str, Any]] = None, ) -> None: """Run a task asynchronously with retries.""" retry_policy = task.retry_policy or retry_policy interval = retry_policy.initial_interval if retry_policy else 0 attempts = 0 config = task.config - if writer is not None: - config = patch_configurable(config, {CONFIG_KEY_SEND: partial(writer, task)}) + if configurable is not None: + config = patch_configurable(config, configurable) while True: try: # clear any writes from previous attempts diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index f46210459..067394bdb 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -1,9 +1,11 @@ import asyncio import concurrent.futures import time +from functools import partial from typing import ( Any, AsyncIterator, + Awaitable, Callable, Iterable, Iterator, @@ -16,6 +18,7 @@ from typing import ( from langgraph.constants import ( CONF, + CONFIG_KEY_CALL, CONFIG_KEY_SEND, ERROR, INTERRUPT, @@ -25,6 +28,7 @@ from langgraph.constants import ( TAG_HIDDEN, ) from langgraph.errors import GraphBubbleUp, GraphInterrupt +from langgraph.pregel.algo import Call from langgraph.pregel.executor import Submit from langgraph.pregel.retry import arun_with_retry, run_with_retry from langgraph.types import PregelExecutableTask, RetryPolicy @@ -41,7 +45,7 @@ class PregelRunner: submit: Submit, put_writes: Callable[[str, Sequence[tuple[str, Any]]], None], schedule_task: Callable[ - [PregelExecutableTask, int], Optional[PregelExecutableTask] + [PregelExecutableTask, int, Optional[Call]], Optional[PregelExecutableTask] ], use_astream: bool = False, node_finished: Optional[Callable[[str], None]] = None, @@ -67,11 +71,10 @@ class PregelRunner: prev_length = len(task.writes) # delegate to the underlying writer task.config[CONF][CONFIG_KEY_SEND](writes) - for idx, w in enumerate(task.writes): - # find the index for the newly inserted writes - if idx < prev_length: - continue - assert writes[idx - prev_length] is w + # confirm no other concurrent writes were added + # TODO could use a lock here instead, if writes can come from many threads + assert len(task.writes) == prev_length + len(writes) + for idx, w in enumerate(writes, start=prev_length): # bail if not a PUSH write if w[0] != PUSH: continue @@ -89,7 +92,10 @@ class PregelRunner: run_with_retry, next_task, retry_policy, - writer=writer, + configurable={ + CONFIG_KEY_SEND: partial(writer, next_task), + # CONFIG_KEY_CALL: partial(call, next_task), + }, __reraise_on_exit__=reraise, ) ] = next_task @@ -102,7 +108,14 @@ class PregelRunner: if len(tasks) == 1 and timeout is None and get_waiter is None: t = tasks[0] try: - run_with_retry(t, retry_policy, writer=writer) + run_with_retry( + t, + retry_policy, + configurable={ + CONFIG_KEY_SEND: partial(writer, t), + # CONFIG_KEY_CALL: partial(call, t), + }, + ) self.commit(t, None) except Exception as exc: self.commit(t, exc) @@ -123,7 +136,10 @@ class PregelRunner: run_with_retry, t, retry_policy, - writer=writer, + configurable={ + CONFIG_KEY_SEND: partial(writer, t), + # CONFIG_KEY_CALL: partial(call, t), + }, __reraise_on_exit__=reraise, ) ] = t @@ -172,21 +188,28 @@ class PregelRunner: get_waiter: Optional[Callable[[], asyncio.Future[None]]] = None, ) -> AsyncIterator[None]: def writer( - task: PregelExecutableTask, writes: Sequence[tuple[str, Any]] - ) -> None: + task: PregelExecutableTask, + writes: Sequence[tuple[str, Any]], + *, + calls: Optional[Sequence[Call]] = None, + ) -> Sequence[Optional[asyncio.Future]]: prev_length = len(task.writes) # delegate to the underlying writer task.config[CONF][CONFIG_KEY_SEND](writes) - for idx, w in enumerate(task.writes): - # find the index for the newly inserted writes - if idx < prev_length: - continue - assert writes[idx - prev_length] is w + # confirm no other concurrent writes were added + # TODO could use a lock here instead, if writes can come from many threads + assert len(task.writes) == prev_length + len(writes) + rtn: dict[int, Optional[asyncio.Future]] = {} + for idx, w in enumerate(writes, start=prev_length): # bail if not a PUSH write if w[0] != PUSH: continue # schedule the next task, if the callback returns one - if next_task := self.schedule_task(task, idx): + if next_task := self.schedule_task( + task, + idx, + calls[idx] if calls is not None else None, + ): # if the parent task was retried, # the next task might already be running if any( @@ -194,21 +217,29 @@ class PregelRunner: ): continue # schedule the next task - futures[ - cast( - asyncio.Future, - self.submit( - arun_with_retry, - next_task, - retry_policy, - stream=self.use_astream, - writer=writer, - __name__=t.name, - __cancel_on_exit__=True, - __reraise_on_exit__=reraise, - ), - ) - ] = next_task + fut = self.submit( + arun_with_retry, + next_task, + retry_policy, + stream=self.use_astream, + configurable={ + CONFIG_KEY_SEND: partial(writer, next_task), + CONFIG_KEY_CALL: partial(call, next_task), + }, + __name__=t.name, + __cancel_on_exit__=True, + __reraise_on_exit__=reraise, + ) + futures[cast(asyncio.Future, fut)] = next_task + rtn[idx] = fut + return [rtn.get(i) for i in range(len(writes))] + + def call( + task, func: str | Callable[[Any], Union[Awaitable[Any], Any]], input: Any + ) -> asyncio.Future[Any]: + (fut,) = writer(task, [(PUSH, None)], calls=[Call(func, input)]) + assert fut is not None, "writer did not return a future for call" + return fut loop = asyncio.get_event_loop() tasks = tuple(tasks) @@ -220,7 +251,13 @@ class PregelRunner: t = tasks[0] try: await arun_with_retry( - t, retry_policy, stream=self.use_astream, writer=writer + t, + retry_policy, + stream=self.use_astream, + configurable={ + CONFIG_KEY_SEND: partial(writer, t), + CONFIG_KEY_CALL: partial(call, t), + }, ) self.commit(t, None) except Exception as exc: @@ -245,7 +282,10 @@ class PregelRunner: t, retry_policy, stream=self.use_astream, - writer=writer, + configurable={ + CONFIG_KEY_SEND: partial(writer, t), + CONFIG_KEY_CALL: partial(call, t), + }, __name__=t.name, __cancel_on_exit__=True, __reraise_on_exit__=reraise, diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 67c7e53f8..b49b1c45e 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -1,9 +1,13 @@ +import asyncio +import concurrent +import concurrent.futures import dataclasses import sys from collections import deque from typing import ( TYPE_CHECKING, Any, + Awaitable, Callable, ClassVar, Generic, @@ -11,6 +15,7 @@ from typing import ( Literal, NamedTuple, Optional, + ParamSpec, Sequence, Type, TypedDict, @@ -363,3 +368,36 @@ def interrupt(value: Any) -> Any: ), ) ) + + +P = ParamSpec("P") +T = TypeVar("T") + + +def call( + func: str | Callable[P, T], *args: P.args, **kwargs: P.kwargs +) -> concurrent.futures.Future[T]: + from langgraph.constants import CONFIG_KEY_CALL + from langgraph.utils.config import get_configurable + + conf = get_configurable() + impl = conf[CONFIG_KEY_CALL] + fut = impl(func, *args, **kwargs) + if not isinstance(fut, concurrent.futures.Future): + raise RuntimeError("In an async context, use acall() instead of call()") + return fut + + +def acall( + func: str | Callable[P, Union[T, Awaitable[T]]], *args: P.args, **kwargs: P.kwargs +) -> asyncio.Future[T]: + from langgraph.constants import CONFIG_KEY_CALL + from langgraph.utils.config import get_configurable + + conf = get_configurable() + impl = conf[CONFIG_KEY_CALL] + fut = impl(func, *args, **kwargs) + if isinstance(fut, concurrent.futures.Future): + fut = asyncio.wrap_future(fut) + fut.cancel + return fut diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index f69d36ed3..120b3deac 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -2496,7 +2496,7 @@ def test_send_react_interrupt( PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()), + path=("__pregel_push", ("__pregel_pull", "agent"), 2), error=None, interrupts=(), state=None, @@ -2653,7 +2653,7 @@ def test_send_react_interrupt( PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()), + path=("__pregel_push", ("__pregel_pull", "agent"), 2), error=None, interrupts=(), state=None, @@ -2740,7 +2740,7 @@ def test_send_react_interrupt( PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", (), 0, AnyStr()), + path=("__pregel_push", (), 0), error=None, interrupts=(), state=None, @@ -2965,7 +2965,7 @@ def test_send_react_interrupt_control( PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()), + path=("__pregel_push", ("__pregel_pull", "agent"), 2), error=None, interrupts=(), state=None, @@ -6043,9 +6043,7 @@ def test_state_graph_packets( ) }, ), - PregelTask( - AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) - ), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), ), next=("tools",), config=(app_w_interrupt.checkpointer.get_tuple(config)).config, @@ -6085,7 +6083,7 @@ def test_state_graph_packets( ), ] }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0, AnyStr())),), + tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0)),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], @@ -6214,12 +6212,8 @@ def test_state_graph_packets( ) }, ), - PregelTask( - AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) - ), - PregelTask( - AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3, AnyStr()) - ), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)), ), next=("tools", "tools"), config=app_w_interrupt.checkpointer.get_tuple(config).config, @@ -6366,9 +6360,7 @@ def test_state_graph_packets( ) }, ), - PregelTask( - AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) - ), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), ), next=("tools",), config=(app_w_interrupt.checkpointer.get_tuple(config)).config, @@ -6408,7 +6400,7 @@ def test_state_graph_packets( ), ] }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0, AnyStr())),), + tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0)),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], @@ -6537,12 +6529,8 @@ def test_state_graph_packets( ) }, ), - PregelTask( - AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) - ), - PregelTask( - AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3, AnyStr()) - ), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)), ), next=("tools", "tools"), config=app_w_interrupt.checkpointer.get_tuple(config).config, @@ -12786,7 +12774,7 @@ def test_send_to_nested_graphs( PregelTask( AnyStr(), "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 1, AnyStr()), + (PUSH, ("__pregel_pull", "__start__"), 1), state={ "configurable": { "thread_id": "1", @@ -12797,7 +12785,7 @@ def test_send_to_nested_graphs( PregelTask( AnyStr(), "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 2, AnyStr()), + (PUSH, ("__pregel_pull", "__start__"), 2), state={ "configurable": { "thread_id": "1", @@ -12850,7 +12838,7 @@ def test_send_to_nested_graphs( "checkpoint_ns": AnyStr("generate_joke:"), "langgraph_checkpoint_ns": AnyStr("generate_joke:"), "langgraph_node": "generate_joke", - "langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 1, AnyStr()], + "langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 1], "langgraph_step": 0, "langgraph_triggers": [PUSH], }, @@ -12895,7 +12883,7 @@ def test_send_to_nested_graphs( "checkpoint_ns": AnyStr("generate_joke:"), "langgraph_checkpoint_ns": AnyStr("generate_joke:"), "langgraph_node": "generate_joke", - "langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 2, AnyStr()], + "langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 2], "langgraph_step": 0, "langgraph_triggers": [PUSH], }, @@ -13021,7 +13009,7 @@ def test_send_to_nested_graphs( PregelTask( AnyStr(), "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 1, AnyStr()), + (PUSH, ("__pregel_pull", "__start__"), 1), state={ "configurable": { "thread_id": "1", @@ -13033,7 +13021,7 @@ def test_send_to_nested_graphs( PregelTask( AnyStr(), "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 2, AnyStr()), + (PUSH, ("__pregel_pull", "__start__"), 2), state={ "configurable": { "thread_id": "1", diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 514703781..d8aeedd9a 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -2864,12 +2864,7 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: PregelTask( id=AnyStr(), name="2", - path=( - "__pregel_push", - ("__pregel_pull", "1"), - 2, - AnyStr(), - ), + path=("__pregel_push", ("__pregel_pull", "1"), 2), error=None, interrupts=(), state=None, @@ -2878,12 +2873,7 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: PregelTask( id=AnyStr(), name="2", - path=( - "__pregel_push", - ("__pregel_pull", "1"), - 3, - AnyStr(), - ), + path=("__pregel_push", ("__pregel_pull", "1"), 3), error=None, interrupts=(), state=None, @@ -2894,14 +2884,8 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: name="2", path=( "__pregel_push", - ( - "__pregel_push", - ("__pregel_pull", "1"), - 2, - AnyStr(), - ), + ("__pregel_push", ("__pregel_pull", "1"), 2), 2, - AnyStr(), ), error=None, interrupts=(), @@ -2913,14 +2897,8 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: name="flaky", path=( "__pregel_push", - ( - "__pregel_push", - ("__pregel_pull", "1"), - 3, - AnyStr(), - ), + ("__pregel_push", ("__pregel_pull", "1"), 3), 2, - AnyStr(), ), error=None, interrupts=(Interrupt(value="Bahh", when="during"),), @@ -3157,7 +3135,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()), + path=("__pregel_push", ("__pregel_pull", "agent"), 2), error=None, interrupts=(), state=None, @@ -3314,7 +3292,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()), + path=("__pregel_push", ("__pregel_pull", "agent"), 2), error=None, interrupts=(), state=None, @@ -3401,7 +3379,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", (), 0, AnyStr()), + path=("__pregel_push", (), 0), error=None, interrupts=(), state=None, @@ -3625,7 +3603,7 @@ async def test_send_react_interrupt_control( PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", ("__pregel_pull", "agent"), 2, AnyStr()), + path=("__pregel_push", ("__pregel_pull", "agent"), 2), error=None, interrupts=(), state=None, @@ -6420,9 +6398,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ) }, ), - PregelTask( - AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) - ), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), ), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, @@ -6465,7 +6441,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ), ] }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0, AnyStr())),), + tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0)),), next=("tools",), config=tup.config, created_at=tup.checkpoint["ts"], @@ -6596,12 +6572,8 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ) }, ), - PregelTask( - AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) - ), - PregelTask( - AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3, AnyStr()) - ), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)), ), next=("tools", "tools"), config=tup.config, @@ -6751,9 +6723,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ) }, ), - PregelTask( - AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) - ), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), ), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, @@ -6796,7 +6766,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ), ] }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0, AnyStr())),), + tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0)),), next=("tools",), config=tup.config, created_at=tup.checkpoint["ts"], @@ -6929,12 +6899,8 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ) }, ), - PregelTask( - AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2, AnyStr()) - ), - PregelTask( - AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3, AnyStr()) - ), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), + PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)), ), next=("tools", "tools"), config=tup.config, @@ -11612,7 +11578,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: PregelTask( AnyStr(), "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 1, AnyStr()), + (PUSH, ("__pregel_pull", "__start__"), 1), state={ "configurable": { "thread_id": "1", @@ -11623,7 +11589,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: PregelTask( AnyStr(), "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 2, AnyStr()), + (PUSH, ("__pregel_pull", "__start__"), 2), state={ "configurable": { "thread_id": "1", @@ -11764,7 +11730,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: PregelTask( AnyStr(), "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 1, AnyStr()), + (PUSH, ("__pregel_pull", "__start__"), 1), state={ "configurable": { "thread_id": "1", @@ -11776,7 +11742,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: PregelTask( AnyStr(), "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 2, AnyStr()), + (PUSH, ("__pregel_pull", "__start__"), 2), state={ "configurable": { "thread_id": "1",