diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index cd847f9be..cb6834f5f 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -1,5 +1,4 @@ import sys -from os import getenv from types import MappingProxyType from typing import Any, Literal, Mapping, cast @@ -93,8 +92,6 @@ 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 NULL_TASK_ID = sys.intern("00000000-0000-0000-0000-000000000000") # the task_id to use for writes that are not associated with a task diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 36c251782..7ae6db6f4 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -966,7 +966,7 @@ class Pregel(PregelProtocol): return patch_checkpoint_map( next_config, saved.metadata if saved else None ) - # no values, copy checkpoint + # no values, empty checkpoint if values is None and as_node is None: next_checkpoint = create_checkpoint(checkpoint, None, step) # copy checkpoint @@ -985,6 +985,7 @@ class Pregel(PregelProtocol): return patch_checkpoint_map( next_config, saved.metadata if saved else None ) + # no values, copy checkpoint if values is None and as_node == "__copy__": next_checkpoint = create_checkpoint(checkpoint, None, step) # copy checkpoint @@ -1248,7 +1249,7 @@ class Pregel(PregelProtocol): return patch_checkpoint_map( next_config, saved.metadata if saved else None ) - # no values, copy checkpoint + # no values, empty checkpoint if values is None and as_node is None: next_checkpoint = create_checkpoint(checkpoint, None, step) # copy checkpoint @@ -1267,6 +1268,7 @@ class Pregel(PregelProtocol): return patch_checkpoint_map( next_config, saved.metadata if saved else None ) + # no values, copy checkpoint if values is None and as_node == "__copy__": next_checkpoint = create_checkpoint(checkpoint, None, step) # copy checkpoint diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 3adea073a..130d7e983 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -228,7 +228,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: t.path[:3]) + tasks = sorted(tasks, key=lambda t: _tuple_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) @@ -273,7 +273,7 @@ def apply_writes( for chan, val in task.writes: if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT, RETURN, ERROR): pass - elif chan == TASKS: # TODO: remove branch in 1.0 + elif chan == TASKS: checkpoint["pending_sends"].append(val) elif chan in channels: pending_writes_by_channel[chan].append(val) @@ -363,8 +363,8 @@ def prepare_next_tasks( This is the union of all PUSH tasks (Sends) and PULL tasks (nodes triggered by edges).""" tasks: list[Union[PregelTask, PregelExecutableTask]] = [] - # Consume pending_sends from previous step (legacy version of Send) - for idx, _ in enumerate(checkpoint["pending_sends"]): # TODO: remove branch in 1.0 + # Consume pending_sends from previous step + for idx, _ in enumerate(checkpoint["pending_sends"]): if task := prepare_single_task( (PUSH, idx), None, @@ -400,65 +400,7 @@ def prepare_next_tasks( manager=manager, ): tasks.append(task) - # Consume pending Sends from this step (new version of Send) - if any(c == PUSH for _, c, _ in pending_writes): - # group writes by task id - grouped_by_task = defaultdict(list) - for tid, c, _ in pending_writes: - grouped_by_task[tid].append(c) - # prepare send tasks from grouped writes - # 1. start from sends originating from existing tasks - tidx = 0 - while tidx < len(tasks): - task = tasks[tidx] - if twrites := grouped_by_task.pop(task.id, None): - for idx, c in enumerate(twrites): - if c != PUSH: - continue - if next_task := prepare_single_task( - (PUSH, task.path, idx, task.id), - None, - checkpoint=checkpoint, - pending_writes=pending_writes, - processes=processes, - channels=channels, - managed=managed, - config=config, - step=step, - for_execution=for_execution, - store=store, - checkpointer=checkpointer, - manager=manager, - ): - tasks.append(next_task) - tidx += 1 - # key tasks by id - task_map = {t.id: t for t in tasks} - # 2. create new tasks for remaining sends (eg. from update_state) - for tid, writes in grouped_by_task.items(): - task = task_map.get(tid) - for idx, c in enumerate(writes): - if c != PUSH: - continue - if next_task := prepare_single_task( - (PUSH, task.path if task else (), idx, tid), - None, - checkpoint=checkpoint, - pending_writes=pending_writes, - processes=processes, - channels=channels, - managed=managed, - config=config, - step=step, - for_execution=for_execution, - store=store, - checkpointer=checkpointer, - manager=manager, - ): - task_map[next_task.id] = next_task - else: - task_map = {t.id: t for t in tasks} - return task_map + return {t.id: t for t in tasks} def prepare_single_task( @@ -571,8 +513,8 @@ def prepare_single_task( else: return PregelTask(task_id, name, task_path[:3]) elif task_path[0] == PUSH: - if len(task_path) == 2: # TODO: remove branch in 1.0 - # legacy SEND tasks, executed in superstep n+1 + if len(task_path) == 2: + # SEND tasks, executed in superstep n+1 # (PUSH, idx of pending send) idx = cast(int, task_path[1]) if idx >= len(checkpoint["pending_sends"]): @@ -601,43 +543,6 @@ def prepare_single_task( PUSH, str(idx), ) - 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_tt = cast(tuple[str, tuple, int, str], task_path) - writes_for_path = [w for w in pending_writes if w[0] == task_path_tt[3]] - if task_path_tt[2] >= len(writes_for_path): - logger.warning( - f"Ignoring invalid write index {task_path[2]} in pending writes" - ) - return - packet = writes_for_path[task_path_tt[2]][2] - if packet is None: - return - if not isinstance(packet, Send): - logger.warning( - f"Ignoring invalid packet type {type(packet)} in pending writes" - ) - return - if packet.node not in processes: - logger.warning( - f"Ignoring unknown node name {packet.node} in pending writes" - ) - return - # create task id - triggers = [PUSH] - checkpoint_ns = ( - f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node - ) - task_id = _uuid5_str( - checkpoint_id, - checkpoint_ns, - str(step), - packet.node, - PUSH, - _tuple_str(task_path[1]), - str(task_path[2]), - ) else: logger.warning(f"Ignoring invalid PUSH task path {task_path}") return @@ -904,7 +809,9 @@ def _uuid5_str(namespace: bytes, *parts: str) -> str: def _tuple_str(tup: Union[str, int, tuple]) -> str: """Generate a string representation of a tuple.""" return ( - f"({', '.join(_tuple_str(x) for x in tup)})" + f"~{', '.join(_tuple_str(x) for x in tup)}" if isinstance(tup, (tuple, list)) + else f"{tup:010d}" + if isinstance(tup, int) else str(tup) ) diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index d54f8b8b2..e7506134b 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -8,10 +8,8 @@ from langgraph.checkpoint.base import PendingWrite from langgraph.constants import ( EMPTY_SEQ, ERROR, - FF_SEND_V2, INTERRUPT, NULL_TASK_ID, - PUSH, RESUME, RETURN, SELF, @@ -82,7 +80,7 @@ def map_command( sends = [cmd.goto] for send in sends: if isinstance(send, Send): - yield (NULL_TASK_ID, PUSH if FF_SEND_V2 else TASKS, send) + yield (NULL_TASK_ID, TASKS, send) elif isinstance(send, str): yield (NULL_TASK_ID, f"branch:{START}:{SELF}:{send}", START) else: diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 89ecc4420..9787c79f6 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -693,9 +693,6 @@ class PregelLoop(LoopProtocol): traceback: Optional[TracebackType], ) -> Optional[bool]: suppress = isinstance(exc_value, GraphInterrupt) and not self.is_nested - if suppress or exc_type is None: - # save final output - self.output = read_channels(self.channels, self.output_keys) if suppress: # emit one last "values" event, with pending writes applied if ( @@ -723,8 +720,13 @@ class PregelLoop(LoopProtocol): "updates", lambda: iter([{INTERRUPT: cast(GraphInterrupt, exc_value).args[0]}]), ) + # save final output + self.output = read_channels(self.channels, self.output_keys) # suppress interrupt return True + elif exc_type is None: + # save final output + self.output = read_channels(self.channels, self.output_keys) def _emit( self, diff --git a/libs/langgraph/langgraph/pregel/write.py b/libs/langgraph/langgraph/pregel/write.py index a4a42c3f0..7469824e4 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, FF_SEND_V2, PUSH, TASKS, Send +from langgraph.constants import CONF, CONFIG_KEY_SEND, TASKS, Send from langgraph.errors import InvalidUpdateError from langgraph.utils.runnable import RunnableCallable @@ -125,7 +125,7 @@ class ChannelWrite(RunnableCallable): # validate for w in writes: if isinstance(w, ChannelWriteEntry): - if w.channel in (TASKS, PUSH): + if w.channel == TASKS: raise InvalidUpdateError( "Cannot write to the reserved channel TASKS" ) @@ -138,7 +138,7 @@ class ChannelWrite(RunnableCallable): tuples: list[tuple[str, Any]] = [] for w in writes: if isinstance(w, Send): - tuples.append((PUSH if FF_SEND_V2 else TASKS, w)) + tuples.append((TASKS, w)) elif isinstance(w, ChannelWriteTupleEntry): if ww := w.mapper(w.value): tuples.extend(ww) diff --git a/libs/langgraph/tests/test_algo.py b/libs/langgraph/tests/test_algo.py index 9d6ec5942..65a2d28c1 100644 --- a/libs/langgraph/tests/test_algo.py +++ b/libs/langgraph/tests/test_algo.py @@ -1,5 +1,6 @@ from langgraph.checkpoint.base import empty_checkpoint -from langgraph.pregel.algo import prepare_next_tasks +from langgraph.constants import PULL, PUSH +from langgraph.pregel.algo import _tuple_str, prepare_next_tasks from langgraph.pregel.manager import ChannelsManager @@ -40,3 +41,26 @@ def test_prepare_next_tasks() -> None: ) # TODO: add more tests + + +def test_tuple_str() -> None: + push_path_a = (PUSH, 2) + pull_path_a = (PULL, "abc") + 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 ( + _tuple_str(push_path_c) + == f"~{PUSH}, ~{PUSH}, ~{PUSH}, 0000000002, 0000000001, 0000000003" + ) + assert _tuple_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)) == [ + f"~{PULL}, abc", + f"~{PUSH}, 0000000002", + f"~{PUSH}, ~{PUSH}, 0000000002, 0000000001", + f"~{PUSH}, ~{PUSH}, ~{PUSH}, 0000000002, 0000000001, 0000000003", + ] diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index f3f399f38..5f0bec3bf 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -16,7 +16,7 @@ from langgraph.channels.context import Context from langgraph.channels.last_value import LastValue from langgraph.channels.untracked_value import UntrackedValue from langgraph.checkpoint.base import BaseCheckpointSaver -from langgraph.constants import END, FF_SEND_V2, PULL, PUSH, START +from langgraph.constants import END, PULL, PUSH, START from langgraph.errors import NodeInterrupt from langgraph.graph import StateGraph from langgraph.graph.graph import Graph @@ -3021,9 +3021,6 @@ def test_state_graph_packets( {"__interrupt__": ()}, ] - if not FF_SEND_V2: - return - assert app_w_interrupt.get_state(config) == StateSnapshot( values={ "messages": [ @@ -3041,33 +3038,7 @@ def test_state_graph_packets( ), ] }, - tasks=( - PregelTask( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ - "messages": AIMessage( - content="", - additional_kwargs={}, - response_metadata={}, - id="ai1", - tool_calls=[ - { - "name": "search_api", - "args": {"query": "query"}, - "id": "tool_call123", - "type": "tool_call", - } - ], - ) - }, - ), - PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), - ), + tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),), next=("tools",), config={ "configurable": { @@ -3080,8 +3051,23 @@ def test_state_graph_packets( metadata={ "parents": {}, "source": "loop", - "step": 0, - "writes": None, + "step": 1, + "writes": { + "agent": { + "messages": AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "query"}, + "id": "tool_call123", + "type": "tool_call", + } + ], + ) + } + }, "thread_id": "1", }, parent_config=( @@ -3116,7 +3102,7 @@ def test_state_graph_packets( ), ] }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0)),), + tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),), next=("tools",), config={ "configurable": { @@ -3129,7 +3115,7 @@ def test_state_graph_packets( metadata={ "parents": {}, "source": "update", - "step": 1, + "step": 2, "writes": { "agent": { "messages": AIMessage( @@ -3227,16 +3213,26 @@ def test_state_graph_packets( ] }, tasks=( - PregelTask( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ + PregelTask(AnyStr(), "tools", (PUSH, 0)), + PregelTask(AnyStr(), "tools", (PUSH, 1)), + ), + next=("tools", "tools"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + created_at=AnyStr(), + metadata={ + "parents": {}, + "source": "loop", + "step": 4, + "writes": { + "agent": { "messages": AIMessage( - "", + content="", id="ai2", tool_calls=[ { @@ -3254,31 +3250,6 @@ def test_state_graph_packets( ], ) }, - ), - PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), - PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)), - ), - next=("tools", "tools"), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={ - "parents": {}, - "source": "loop", - "step": 2, - "writes": { - "tools": { - "messages": _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ), - }, }, "thread_id": "1", }, @@ -3334,7 +3305,7 @@ def test_state_graph_packets( metadata={ "parents": {}, "source": "update", - "step": 3, + "step": 5, "writes": { "agent": { "messages": AIMessage(content="answer", id="ai2"), @@ -3400,31 +3371,7 @@ def test_state_graph_packets( ), ] }, - tasks=( - PregelTask( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ - "messages": AIMessage( - "", - id="ai1", - tool_calls=[ - { - "name": "search_api", - "args": {"query": "query"}, - "id": "tool_call123", - "type": "tool_call", - } - ], - ) - }, - ), - PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), - ), + tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),), next=("tools",), config={ "configurable": { @@ -3437,8 +3384,23 @@ def test_state_graph_packets( metadata={ "parents": {}, "source": "loop", - "step": 0, - "writes": None, + "step": 1, + "writes": { + "agent": { + "messages": AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "query"}, + "id": "tool_call123", + "type": "tool_call", + } + ], + ) + } + }, "thread_id": "2", }, parent_config=( @@ -3473,14 +3435,14 @@ def test_state_graph_packets( ), ] }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0)),), + tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=AnyStr(), metadata={ "parents": {}, "source": "update", - "step": 1, + "step": 2, "writes": { "agent": { "messages": AIMessage( @@ -3578,36 +3540,8 @@ def test_state_graph_packets( ] }, tasks=( - PregelTask( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ - "messages": AIMessage( - "", - id="ai2", - tool_calls=[ - { - "name": "search_api", - "args": {"query": "another", "idx": 0}, - "id": "tool_call234", - "type": "tool_call", - }, - { - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - "id": "tool_call567", - "type": "tool_call", - }, - ], - ) - }, - ), - PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), - PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)), + PregelTask(AnyStr(), "tools", (PUSH, 0)), + PregelTask(AnyStr(), "tools", (PUSH, 1)), ), next=("tools", "tools"), config={ @@ -3621,14 +3555,25 @@ def test_state_graph_packets( metadata={ "parents": {}, "source": "loop", - "step": 2, + "step": 4, "writes": { - "tools": { - "messages": _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", - ), + "agent": { + "messages": AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], + ) }, }, "thread_id": "2", @@ -3685,7 +3630,7 @@ def test_state_graph_packets( metadata={ "parents": {}, "source": "update", - "step": 3, + "step": 5, "writes": { "agent": { "messages": AIMessage(content="answer", id="ai2"), @@ -5859,7 +5804,6 @@ def test_dynamic_interrupt( ) -@pytest.mark.skipif(not FF_SEND_V2, reason="send v2 is not enabled") @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_copy_checkpoint( request: pytest.FixtureRequest, checkpointer_name: str @@ -5879,6 +5823,7 @@ def test_copy_checkpoint( nonlocal tool_two_node_count tool_two_node_count += 1 if s["market"] == "DE": + time.sleep(0.1) answer = interrupt("Just because...") else: answer = " all good" @@ -5908,7 +5853,7 @@ def test_copy_checkpoint( assert run.outputs == {"market": "DE", "my_key": "value one"} assert tool_two.invoke({"my_key": "value", "market": "US"}) == { - "my_key": "value one all good", + "my_key": "value all good one", "market": "US", } @@ -5939,6 +5884,7 @@ def test_copy_checkpoint( ] # resume with answer assert [c for c in tool_two.stream(Command(resume=" my answer"), thread2)] == [ + {"tool_one": {"my_key": " one"}, "__metadata__": {"cached": True}}, {"tool_two": {"my_key": " my answer"}}, ] @@ -5955,7 +5901,7 @@ def test_copy_checkpoint( "parents": {}, "source": "loop", "step": 0, - "writes": {"tool_one": {"my_key": " one"}}, + "writes": None, "thread_id": "1", }, { @@ -5971,6 +5917,12 @@ def test_copy_checkpoint( values={"my_key": "value ⛰️ one", "market": "DE"}, next=("tool_two",), tasks=( + PregelTask( + id=AnyStr(), + name="tool_one", + path=("__pregel_push", 0), + result={"my_key": " one"}, + ), PregelTask( AnyStr(), "tool_two", @@ -5996,7 +5948,7 @@ def test_copy_checkpoint( "parents": {}, "source": "loop", "step": 0, - "writes": {"tool_one": {"my_key": " one"}}, + "writes": None, "thread_id": "1", }, parent_config=( @@ -6005,13 +5957,25 @@ def test_copy_checkpoint( else [*tool_two.checkpointer.list(thread1, limit=2)][-1].config ), ) + + if "shallow" in checkpointer_name: + return + # clear the interrupt and next tasks - tool_two.update_state(thread1, None) + tool_two.update_state(thread1, None, as_node="__copy__") # interrupt is cleared, next task is kept assert tool_two.get_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️ one", "market": "DE"}, - next=("tool_two",), + values={"my_key": "value ⛰️", "market": "DE"}, + next=( + "tool_one", + "tool_two", + ), tasks=( + PregelTask( + id=AnyStr(), + name="tool_one", + path=("__pregel_push", 0), + ), PregelTask( AnyStr(), "tool_two", @@ -6029,15 +5993,13 @@ def test_copy_checkpoint( created_at=AnyStr(), metadata={ "parents": {}, - "source": "update", + "source": "fork", "step": 1, - "writes": {}, + "writes": None, "thread_id": "1", }, parent_config=( - None - if "shallow" in checkpointer_name - else [*tool_two.checkpointer.list(thread1, limit=2)][-1].config + [*tool_two.checkpointer.list(thread1, limit=2)][-1].parent_config ), ) @@ -7292,12 +7254,11 @@ 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 ) -> None: - if not FF_SEND_V2: - pytest.skip("Send deduplication is only available in Send V2") checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") class InterruptOnce: @@ -9320,328 +9281,18 @@ 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=( - PregelTask( - id=AnyStr(), - name="__start__", - path=("__pregel_pull", "__start__"), - error=None, - interrupts=(), - state=None, - result={"subjects": ["cats", "dogs"]}, - ), - PregelTask( - AnyStr(), - "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 1), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, - ), - PregelTask( - AnyStr(), - "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 2), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, - ), - ), - next=("generate_joke", "generate_joke"), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"__start__": {"subjects": ["cats", "dogs"]}}, - "step": -1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - ) - # check state of each of the inner tasks - assert graph.get_state(outer_state.tasks[1].state) == StateSnapshot( - values={"subject": "cats - hohoho", "jokes": []}, - next=("generate",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("generate_joke:"): AnyStr(), - } - ), - } - }, - metadata={ - "step": 1, - "source": "loop", - "writes": None, - "parents": {"": AnyStr()}, - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - "langgraph_checkpoint_ns": AnyStr("generate_joke:"), - "langgraph_node": "generate_joke", - "langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 1], - "langgraph_step": 0, - "langgraph_triggers": [PUSH], - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("generate_joke:"): AnyStr(), - } - ), - } - } - ), - tasks=(PregelTask(id=AnyStr(""), name="generate", path=(PULL, "generate")),), - ) - assert graph.get_state(outer_state.tasks[2].state) == StateSnapshot( - values={"subject": "dogs - hohoho", "jokes": []}, - next=("generate",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("generate_joke:"): AnyStr(), - } - ), - } - }, - metadata={ - "step": 1, - "source": "loop", - "writes": None, - "parents": {"": AnyStr()}, - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - "langgraph_checkpoint_ns": AnyStr("generate_joke:"), - "langgraph_node": "generate_joke", - "langgraph_path": [PUSH, ["__pregel_pull", "__start__"], 2], - "langgraph_step": 0, - "langgraph_triggers": [PUSH], - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("generate_joke:"): AnyStr(), - } - ), - } - } - ), - tasks=(PregelTask(id=AnyStr(""), name="generate", path=(PULL, "generate")),), - ) # update state of dogs joke graph - graph.update_state( - outer_state.tasks[2 if FF_SEND_V2 else 1].state, {"subject": "turtles - hohoho"} - ) + 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] + 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"]}}, ] - actual_snapshot = graph.get_state(config) - expected_snapshot = StateSnapshot( - values={ - "subjects": ["cats", "dogs"], - "jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"], - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "generate_joke": [ - {"jokes": ["Joke about cats - hohoho"]}, - {"jokes": ["Joke about turtles - hohoho"]}, - ] - }, - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - ) - assert actual_snapshot == expected_snapshot - - if "shallow" in checkpointer_name: - return - - # test full history - actual_history = list(graph.get_state_history(config)) - - # get subgraph node state for expected history - expected_history = [ - StateSnapshot( - values={ - "subjects": ["cats", "dogs"], - "jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"], - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "generate_joke": [ - {"jokes": ["Joke about cats - hohoho"]}, - {"jokes": ["Joke about turtles - hohoho"]}, - ] - }, - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"jokes": []}, - tasks=( - PregelTask( - id=AnyStr(), - name="__start__", - path=("__pregel_pull", "__start__"), - error=None, - interrupts=(), - state=None, - result={"subjects": ["cats", "dogs"]}, - ), - PregelTask( - AnyStr(), - "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 1), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, - result={"jokes": ["Joke about cats - hohoho"]}, - ), - PregelTask( - AnyStr(), - "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 2), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, - result={"jokes": ["Joke about turtles - hohoho"]}, - ), - ), - next=("__start__", "generate_joke", "generate_joke"), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"__start__": {"subjects": ["cats", "dogs"]}}, - "step": -1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - assert actual_history == expected_history - @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_send_react_interrupt( @@ -9767,9 +9418,6 @@ 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( @@ -9798,9 +9446,24 @@ def test_send_react_interrupt( } }, metadata={ - "step": 0, + "step": 1, "source": "loop", - "writes": None, + "writes": { + "agent": { + "messages": AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ) + } + }, "parents": {}, "thread_id": "2", }, @@ -9817,34 +9480,10 @@ def test_send_react_interrupt( } ), tasks=( - PregelTask( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ - "messages": AIMessage( - content="", - additional_kwargs={}, - response_metadata={}, - id="ai1", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ) - }, - ), PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", ("__pregel_pull", "agent"), 2), + path=("__pregel_push", 0), error=None, interrupts=(), state=None, @@ -9878,7 +9517,7 @@ def test_send_react_interrupt( } }, metadata={ - "step": 1, + "step": 2, "source": "update", "writes": { "agent": { @@ -9965,9 +9604,26 @@ def test_send_react_interrupt( } }, metadata={ - "step": 0, + "step": 1, "source": "loop", - "writes": None, + "writes": { + "agent": { + "messages": AIMessage( + content="", + additional_kwargs={}, + response_metadata={}, + id="ai1", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ) + } + }, "parents": {}, "thread_id": "3", }, @@ -9984,32 +9640,10 @@ def test_send_react_interrupt( } ), tasks=( - PregelTask( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ - "messages": AIMessage( - "", - id="ai1", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ) - }, - ), PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", ("__pregel_pull", "agent"), 2), + path=("__pregel_push", 0), error=None, interrupts=(), state=None, @@ -10064,7 +9698,7 @@ def test_send_react_interrupt( } }, metadata={ - "step": 1, + "step": 2, "source": "update", "writes": { "agent": { @@ -10100,7 +9734,7 @@ def test_send_react_interrupt( PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", (), 0), + path=("__pregel_push", 0), error=None, interrupts=(), state=None, @@ -10230,9 +9864,6 @@ 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"]) @@ -10283,9 +9914,24 @@ def test_send_react_interrupt_control( } }, metadata={ - "step": 0, + "step": 1, "source": "loop", - "writes": None, + "writes": { + "agent": { + "messages": AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ) + } + }, "parents": {}, "thread_id": "2", }, @@ -10302,34 +9948,10 @@ def test_send_react_interrupt_control( } ), tasks=( - PregelTask( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ - "messages": AIMessage( - content="", - additional_kwargs={}, - response_metadata={}, - id="ai1", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ) - }, - ), PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", ("__pregel_pull", "agent"), 2), + path=("__pregel_push", 0), error=None, interrupts=(), state=None, @@ -10363,7 +9985,7 @@ def test_send_react_interrupt_control( } }, metadata={ - "step": 1, + "step": 2, "source": "update", "writes": { "agent": { diff --git a/libs/langgraph/tests/test_large_cases_async.py b/libs/langgraph/tests/test_large_cases_async.py index f6e807fca..dab4f3f35 100644 --- a/libs/langgraph/tests/test_large_cases_async.py +++ b/libs/langgraph/tests/test_large_cases_async.py @@ -25,7 +25,7 @@ from syrupy import SnapshotAssertion from langgraph.channels.context import Context from langgraph.channels.last_value import LastValue from langgraph.channels.untracked_value import UntrackedValue -from langgraph.constants import END, FF_SEND_V2, PULL, PUSH, START +from langgraph.constants import END, PULL, PUSH, START from langgraph.graph.graph import Graph from langgraph.graph.message import MessageGraph, add_messages from langgraph.graph.state import StateGraph @@ -2740,9 +2740,6 @@ 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": [ @@ -2760,31 +2757,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ), ] }, - tasks=( - PregelTask( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ - "messages": AIMessage( - "", - id="ai1", - tool_calls=[ - { - "name": "search_api", - "args": {"query": "query"}, - "id": "tool_call123", - "type": "tool_call", - } - ], - ) - }, - ), - PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), - ), + tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),), next=("tools",), config={ "configurable": { @@ -2797,8 +2770,23 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: metadata={ "parents": {}, "source": "loop", - "step": 0, - "writes": None, + "step": 1, + "writes": { + "agent": { + "messages": AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "name": "search_api", + "args": {"query": "query"}, + "id": "tool_call123", + "type": "tool_call", + } + ], + ) + } + }, "thread_id": "1", }, parent_config=( @@ -2834,14 +2822,14 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ), ] }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0)),), + tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),), next=("tools",), config=tup.config, created_at=tup.checkpoint["ts"], metadata={ "parents": {}, "source": "update", - "step": 1, + "step": 2, "writes": { "agent": { "messages": AIMessage( @@ -2941,36 +2929,8 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ] }, tasks=( - PregelTask( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ - "messages": AIMessage( - "", - id="ai2", - tool_calls=[ - { - "name": "search_api", - "args": {"query": "another", "idx": 0}, - "id": "tool_call234", - "type": "tool_call", - }, - { - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - "id": "tool_call567", - "type": "tool_call", - }, - ], - ) - }, - ), - PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), - PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)), + PregelTask(AnyStr(), "tools", (PUSH, 0)), + PregelTask(AnyStr(), "tools", (PUSH, 1)), ), next=("tools", "tools"), config=tup.config, @@ -2978,13 +2938,24 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: metadata={ "parents": {}, "source": "loop", - "step": 2, + "step": 4, "writes": { - "tools": { - "messages": _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", + "agent": { + "messages": AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], ), }, }, @@ -3036,7 +3007,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: metadata={ "parents": {}, "source": "update", - "step": 3, + "step": 5, "writes": { "agent": { "messages": AIMessage(content="answer", id="ai2"), @@ -3103,15 +3074,16 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ), ] }, - tasks=( - PregelTask( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ + tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),), + next=("tools",), + config=tup.config, + created_at=tup.checkpoint["ts"], + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "writes": { + "agent": { "messages": AIMessage( content="", additional_kwargs={}, @@ -3126,18 +3098,8 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: } ], ) - }, - ), - PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), - ), - next=("tools",), - config=tup.config, - created_at=tup.checkpoint["ts"], - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "writes": None, + } + }, "thread_id": "2", }, parent_config=( @@ -3173,14 +3135,14 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ), ] }, - tasks=(PregelTask(AnyStr(), "tools", (PUSH, (), 0)),), + tasks=(PregelTask(AnyStr(), "tools", (PUSH, 0)),), next=("tools",), config=tup.config, created_at=tup.checkpoint["ts"], metadata={ "parents": {}, "source": "update", - "step": 1, + "step": 2, "writes": { "agent": { "messages": AIMessage( @@ -3280,38 +3242,8 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: ] }, tasks=( - PregelTask( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ - "messages": AIMessage( - content="", - additional_kwargs={}, - response_metadata={}, - id="ai2", - tool_calls=[ - { - "name": "search_api", - "args": {"query": "another", "idx": 0}, - "id": "tool_call234", - "type": "tool_call", - }, - { - "name": "search_api", - "args": {"query": "a third one", "idx": 1}, - "id": "tool_call567", - "type": "tool_call", - }, - ], - ) - }, - ), - PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 2)), - PregelTask(AnyStr(), "tools", (PUSH, ("__pregel_pull", "agent"), 3)), + PregelTask(AnyStr(), "tools", (PUSH, 0)), + PregelTask(AnyStr(), "tools", (PUSH, 1)), ), next=("tools", "tools"), config=tup.config, @@ -3319,13 +3251,24 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: metadata={ "parents": {}, "source": "loop", - "step": 2, + "step": 4, "writes": { - "tools": { - "messages": _AnyIdToolMessage( - content="result for a different query", - name="search_api", - tool_call_id="tool_call123", + "agent": { + "messages": AIMessage( + id="ai2", + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another", "idx": 0}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one", "idx": 1}, + }, + ], ), }, }, @@ -3377,7 +3320,7 @@ async def test_state_graph_packets(checkpointer_name: str) -> None: metadata={ "parents": {}, "source": "update", - "step": 3, + "step": 5, "writes": { "agent": { "messages": AIMessage(content="answer", id="ai2"), @@ -6961,83 +6904,9 @@ 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=( - PregelTask( - id=AnyStr(), - name="__start__", - path=("__pregel_pull", "__start__"), - error=None, - interrupts=(), - state=None, - result={"subjects": ["cats", "dogs"]}, - ), - PregelTask( - AnyStr(), - "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 1), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, - ), - PregelTask( - AnyStr(), - "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 2), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, - ), - ), - next=("generate_joke", "generate_joke"), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": { - "__start__": { - "subjects": [ - "cats", - "dogs", - ], - } - }, - "step": -1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - ) - # update state of dogs joke graph await graph.aupdate_state( - outer_state.tasks[2].state, {"subject": "turtles - hohoho"} + outer_state.tasks[1].state, {"subject": "turtles - hohoho"} ) # continue past interrupt @@ -7046,150 +6915,6 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None: "jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"], } - actual_snapshot = await graph.aget_state(config) - expected_snapshot = StateSnapshot( - values={ - "subjects": ["cats", "dogs"], - "jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"], - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "generate_joke": [ - {"jokes": ["Joke about cats - hohoho"]}, - {"jokes": ["Joke about turtles - hohoho"]}, - ] - }, - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=( - None - if "shallow" in checkpointer_name - else { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - ) - assert actual_snapshot == expected_snapshot - - if "shallow" in checkpointer_name: - return - - # test full history - actual_history = [c async for c in graph.aget_state_history(config)] - expected_history = [ - StateSnapshot( - values={ - "subjects": ["cats", "dogs"], - "jokes": [ - "Joke about cats - hohoho", - "Joke about turtles - hohoho", - ], - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "generate_joke": [ - {"jokes": ["Joke about cats - hohoho"]}, - {"jokes": ["Joke about turtles - hohoho"]}, - ] - }, - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"jokes": []}, - next=("__start__", "generate_joke", "generate_joke"), - tasks=( - PregelTask( - id=AnyStr(), - name="__start__", - path=("__pregel_pull", "__start__"), - error=None, - interrupts=(), - state=None, - result={"subjects": ["cats", "dogs"]}, - ), - PregelTask( - AnyStr(), - "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 1), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, - result={"jokes": ["Joke about cats - hohoho"]}, - ), - PregelTask( - AnyStr(), - "generate_joke", - (PUSH, ("__pregel_pull", "__start__"), 2), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("generate_joke:"), - } - }, - result={"jokes": ["Joke about turtles - hohoho"]}, - ), - ), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"__start__": {"subjects": ["cats", "dogs"]}}, - "step": -1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - assert actual_history == expected_history - @pytest.mark.skipif( sys.version_info < (3, 11), diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index b08cc87a7..88e748fae 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -50,13 +50,7 @@ from langgraph.checkpoint.base import ( CheckpointTuple, ) from langgraph.checkpoint.memory import MemorySaver -from langgraph.constants import ( - CONFIG_KEY_NODE_FINISHED, - ERROR, - FF_SEND_V2, - PULL, - START, -) +from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, START from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError from langgraph.func import entrypoint, task from langgraph.graph import END, Graph, StateGraph @@ -1375,31 +1369,17 @@ 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", - ] - if FF_SEND_V2 - else [ - "0", - "1", - "1.1", - "3.1", - "2|1", - "2|2", - "2|3", - "2|4", - "3", - ] - ) + assert graph.invoke(["0"]) == [ + "0", + "1", + "1.1", + "3.1", + "2|1", + "2|2", + "2|3", + "2|4", + "3", + ] def test_send_sequences() -> None: @@ -1438,31 +1418,17 @@ 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(goto=Send(node='2', arg=3))", - "2|Command(goto=Send(node='2', arg=4))", - "2|3", - "2|4", - "3", - "3.1", - ] - if FF_SEND_V2 - else [ - "0", - "1", - "3.1", - "2|Command(goto=Send(node='2', arg=3))", - "2|Command(goto=Send(node='2', arg=4))", - "3", - "2|3", - "2|4", - "3", - ] - ) + assert graph.invoke(["0"]) == [ + "0", + "1", + "3.1", + "2|Command(goto=Send(node='2', arg=3))", + "2|Command(goto=Send(node='2', arg=4))", + "3", + "2|3", + "2|4", + "3", + ] @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index bb3e9ba08..347e31574 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -47,13 +47,7 @@ from langgraph.checkpoint.base import ( CheckpointTuple, ) from langgraph.checkpoint.memory import MemorySaver -from langgraph.constants import ( - CONFIG_KEY_NODE_FINISHED, - ERROR, - FF_SEND_V2, - PULL, - START, -) +from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH, START from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt from langgraph.func import entrypoint, task from langgraph.graph import END, Graph, StateGraph @@ -878,7 +872,6 @@ async def test_dynamic_interrupt_subgraph(checkpointer_name: str) -> None: ) -@pytest.mark.skipif(not FF_SEND_V2, reason="send v2 is not enabled") @pytest.mark.skipif( sys.version_info < (3, 11), reason="Python 3.11+ is required for async contextvars support", @@ -914,7 +907,7 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None: tracer = FakeTracer() assert await tool_two.ainvoke( - {"my_key": "value", "market": "DE"}, {"callbacks": [tracer]} + {"my_key": "value", "market": "DE"}, {"callbacks": [tracer]}, debug=True ) == { "my_key": "value one", "market": "DE", @@ -927,7 +920,7 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None: assert run.outputs == {"market": "DE", "my_key": "value one"} assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == { - "my_key": "value one all good", + "my_key": "value all good one", "market": "US", } @@ -964,6 +957,10 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None: assert [ c async for c in tool_two.astream(Command(resume=" my answer"), thread2) ] == [ + { + "__metadata__": {"cached": True}, + "tool_one": {"my_key": " one"}, + }, {"tool_two": {"my_key": " my answer"}}, ] @@ -983,7 +980,7 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None: "parents": {}, "source": "loop", "step": 0, - "writes": {"tool_one": {"my_key": " one"}}, + "writes": None, "thread_id": "1", }, { @@ -1000,6 +997,15 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None: values={"my_key": "value ⛰️ one", "market": "DE"}, next=("tool_two",), tasks=( + PregelTask( + AnyStr(), + name="tool_one", + path=("__pregel_push", 0), + error=None, + interrupts=(), + state=None, + result={"my_key": " one"}, + ), PregelTask( AnyStr(), "tool_two", @@ -1019,7 +1025,7 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None: "parents": {}, "source": "loop", "step": 0, - "writes": {"tool_one": {"my_key": " one"}}, + "writes": None, "thread_id": "1", }, parent_config=( @@ -1030,14 +1036,25 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None: ].config ), ) + + if "shallow" in checkpointer_name: + # shallow checkpointer doesn't support copy + return + # clear the interrupt and next tasks - await tool_two.aupdate_state(thread1, None) + await tool_two.aupdate_state(thread1, None, as_node="__copy__") # interrupt is cleared, next task is kept tup = await tool_two.checkpointer.aget_tuple(thread1) assert await tool_two.aget_state(thread1) == StateSnapshot( - values={"my_key": "value ⛰️ one", "market": "DE"}, - next=("tool_two",), + values={"my_key": "value ⛰️", "market": "DE"}, + next=("tool_one", "tool_two"), tasks=( + PregelTask( + AnyStr(), + "tool_one", + (PUSH, 0), + result=None, + ), PregelTask( AnyStr(), "tool_two", @@ -1049,17 +1066,15 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None: created_at=tup.checkpoint["ts"], metadata={ "parents": {}, - "source": "update", + "source": "fork", "step": 1, - "writes": {}, + "writes": None, "thread_id": "1", }, parent_config=( - None - if "shallow" in checkpointer_name - else [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ + [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ -1 - ].config + ].parent_config ), ) @@ -2345,18 +2360,6 @@ async def test_concurrent_emit_sends() -> None: graph = builder.compile() 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", @@ -2370,6 +2373,7 @@ 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: @@ -2407,34 +2411,17 @@ 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(goto=Send(node='2', arg=3))", - "2|Command(goto=Send(node='2', arg=4))", - "2|3", - "2|4", - "3", - "3.1", - ] - if FF_SEND_V2 - else [ - "0", - "1", - "3.1", - "2|Command(goto=Send(node='2', arg=3))", - "2|Command(goto=Send(node='2', arg=4))", - "3", - "2|3", - "2|4", - "3", - ] - ) - - if not FF_SEND_V2: - return + assert await graph.ainvoke(["0"]) == [ + "0", + "1", + "3.1", + "2|Command(goto=Send(node='2', arg=3))", + "2|Command(goto=Send(node='2', arg=4))", + "3", + "2|3", + "2|4", + "3", + ] async with awith_checkpointer(checkpointer_name) as checkpointer: graph = builder.compile(checkpointer=checkpointer, interrupt_before=["3.1"]) @@ -2442,20 +2429,17 @@ async def test_send_sequences(checkpointer_name: str) -> None: assert await graph.ainvoke(["0"], thread1) == [ "0", "1", - "2|Command(goto=Send(node='2', arg=3))", - "2|Command(goto=Send(node='2', arg=4))", - "2|3", - "2|4", ] assert await graph.ainvoke(None, thread1) == [ "0", "1", + "3.1", "2|Command(goto=Send(node='2', arg=3))", "2|Command(goto=Send(node='2', arg=4))", + "3", "2|3", "2|4", "3", - "3.1", ] @@ -2631,11 +2615,9 @@ 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: - if not FF_SEND_V2: - pytest.skip("Send deduplication is only available in Send V2") - class InterruptOnce: ticks: int = 0 @@ -2689,22 +2671,26 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: assert await graph.ainvoke(["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 + print((await graph.aget_state(thread1)).tasks) # resume execution assert await graph.ainvoke(None, thread1, debug=1) == [ "0", "1", + "3.1", "2|Command(goto=Send(node='2', arg=3))", "2|Command(goto=Send(node='flaky', arg=4))", - "2|3", - "flaky|4", "3", - "3.1", + "flaky|4", + "2|3", + "3", ] # node "2" doesn't get called again, as we recover writes saved before assert builder.nodes["2"].runnable.func.ticks == 3 @@ -2717,12 +2703,13 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: values=[ "0", "1", + "3.1", "2|Command(goto=Send(node='2', arg=3))", "2|Command(goto=Send(node='flaky', arg=4))", - "2|3", - "flaky|4", "3", - "3.1", + "flaky|4", + "2|3", + "3", ], next=(), config={ @@ -2734,9 +2721,9 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: }, metadata={ "source": "loop", - "writes": {"3": ["3"], "3.1": ["3.1"]}, + "writes": {"3": ["3"]}, "thread_id": "1", - "step": 2, + "step": 4, "parents": {}, }, created_at=AnyStr(), @@ -2753,12 +2740,14 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: values=[ "0", "1", + "3.1", "2|Command(goto=Send(node='2', arg=3))", "2|Command(goto=Send(node='flaky', arg=4))", - "2|3", + "3", "flaky|4", + "2|3", ], - next=("3", "3.1"), + next=("3",), config={ "configurable": { "thread_id": "1", @@ -2768,17 +2757,9 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: }, 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(), @@ -2799,6 +2780,123 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: 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", @@ -2812,7 +2910,7 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: ), StateSnapshot( values=["0"], - next=("1", "2", "2", "2", "flaky"), + next=("1",), config={ "configurable": { "thread_id": "1", @@ -2845,50 +2943,6 @@ async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: 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( @@ -3046,9 +3100,6 @@ 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( @@ -3077,9 +3128,24 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: } }, metadata={ - "step": 0, + "step": 1, "source": "loop", - "writes": None, + "writes": { + "agent": { + "messages": AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ) + } + }, "parents": {}, "thread_id": "2", }, @@ -3096,34 +3162,10 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: } ), tasks=( - PregelTask( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ - "messages": AIMessage( - content="", - additional_kwargs={}, - response_metadata={}, - id="ai1", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ) - }, - ), PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", ("__pregel_pull", "agent"), 2), + path=("__pregel_push", 0), error=None, interrupts=(), state=None, @@ -3157,7 +3199,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: } }, metadata={ - "step": 1, + "step": 2, "source": "update", "writes": { "agent": { @@ -3244,9 +3286,24 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: } }, metadata={ - "step": 0, + "step": 1, "source": "loop", - "writes": None, + "writes": { + "agent": { + "messages": AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ) + } + }, "parents": {}, "thread_id": "3", }, @@ -3263,32 +3320,10 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: } ), tasks=( - PregelTask( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ - "messages": AIMessage( - "", - id="ai1", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ) - }, - ), PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", ("__pregel_pull", "agent"), 2), + path=("__pregel_push", 0), error=None, interrupts=(), state=None, @@ -3343,7 +3378,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: } }, metadata={ - "step": 1, + "step": 2, "source": "update", "writes": { "agent": { @@ -3379,7 +3414,7 @@ async def test_send_react_interrupt(checkpointer_name: str) -> None: PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", (), 0), + path=("__pregel_push", 0), error=None, interrupts=(), state=None, @@ -3530,9 +3565,6 @@ async def test_send_react_interrupt_control( } 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( @@ -3561,9 +3593,24 @@ async def test_send_react_interrupt_control( } }, metadata={ - "step": 0, + "step": 1, "source": "loop", - "writes": None, + "writes": { + "agent": { + "messages": AIMessage( + content="", + id="ai1", + tool_calls=[ + { + "name": "foo", + "args": {"hi": [1, 2, 3]}, + "id": "", + "type": "tool_call", + } + ], + ) + } + }, "parents": {}, "thread_id": "2", }, @@ -3580,34 +3627,10 @@ async def test_send_react_interrupt_control( } ), tasks=( - PregelTask( - id=AnyStr(), - name="agent", - path=("__pregel_pull", "agent"), - error=None, - interrupts=(), - state=None, - result={ - "messages": AIMessage( - content="", - additional_kwargs={}, - response_metadata={}, - id="ai1", - tool_calls=[ - { - "name": "foo", - "args": {"hi": [1, 2, 3]}, - "id": "", - "type": "tool_call", - } - ], - ) - }, - ), PregelTask( id=AnyStr(), name="foo", - path=("__pregel_push", ("__pregel_pull", "agent"), 2), + path=("__pregel_push", 0), error=None, interrupts=(), state=None, @@ -3641,7 +3664,7 @@ async def test_send_react_interrupt_control( } }, metadata={ - "step": 1, + "step": 2, "source": "update", "writes": { "agent": { @@ -5732,9 +5755,10 @@ async def test_store_injected_async(checkpointer_name: str, store_name: str) -> builder.add_node(f"node_{i}", Node(i)) builder.add_edge("__start__", f"node_{i}") - async with awith_checkpointer(checkpointer_name) as checkpointer, awith_store( - store_name - ) as the_store: + async with ( + awith_checkpointer(checkpointer_name) as checkpointer, + awith_store(store_name) as the_store, + ): graph = builder.compile(store=the_store, checkpointer=checkpointer) # Test batch operations with multiple threads diff --git a/libs/scheduler-kafka/tests/test_push.py b/libs/scheduler-kafka/tests/test_push.py index 3d2e4d43d..7bebc7d5a 100644 --- a/libs/scheduler-kafka/tests/test_push.py +++ b/libs/scheduler-kafka/tests/test_push.py @@ -9,7 +9,7 @@ import pytest from aiokafka import AIOKafkaProducer from langgraph.checkpoint.base import BaseCheckpointSaver -from langgraph.constants import FF_SEND_V2, START +from langgraph.constants import START from langgraph.errors import NodeInterrupt from langgraph.graph.state import CompiledStateGraph, StateGraph from langgraph.scheduler.kafka import serde @@ -76,10 +76,8 @@ def mk_push_graph( return builder.compile(checkpointer=checkpointer) +@pytest.mark.skip("TODO: re-enable in next PR") async def test_push_graph(topics: Topics, acheckpointer: BaseCheckpointSaver) -> None: - if not FF_SEND_V2: - pytest.skip("Test requires FF_SEND_V2") - input = ["0"] config = {"configurable": {"thread_id": "1"}} graph = mk_push_graph(acheckpointer) diff --git a/libs/scheduler-kafka/tests/test_push_sync.py b/libs/scheduler-kafka/tests/test_push_sync.py index ee33d613e..ea9fc6463 100644 --- a/libs/scheduler-kafka/tests/test_push_sync.py +++ b/libs/scheduler-kafka/tests/test_push_sync.py @@ -8,7 +8,7 @@ from typing import ( import pytest from langgraph.checkpoint.base import BaseCheckpointSaver -from langgraph.constants import FF_SEND_V2, START +from langgraph.constants import START from langgraph.errors import NodeInterrupt from langgraph.graph.state import CompiledStateGraph, StateGraph from langgraph.scheduler.kafka import serde @@ -76,10 +76,8 @@ def mk_push_graph( return builder.compile(checkpointer=checkpointer) +@pytest.mark.skip("TODO: re-enable in next PR") def test_push_graph(topics: Topics, acheckpointer: BaseCheckpointSaver) -> None: - if not FF_SEND_V2: - pytest.skip("Test requires FF_SEND_V2") - input = ["0"] config = {"configurable": {"thread_id": "1"}} graph = mk_push_graph(acheckpointer)