From b3ee7288396d58ea7d18c74ef793c2c7b8b2d136 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 18 Jul 2024 15:48:28 -0700 Subject: [PATCH] Pregel.stream passing all tests --- libs/langgraph/langgraph/pregel/__init__.py | 88 +++++------- libs/langgraph/langgraph/pregel/algo.py | 1 - libs/langgraph/langgraph/pregel/debug.py | 4 + libs/langgraph/langgraph/pregel/loop.py | 149 ++++++++++++-------- libs/langgraph/tests/test_prebuilt.py | 118 +++++++++++++++- libs/langgraph/tests/test_pregel.py | 19 ++- libs/langgraph/tests/test_pregel_async.py | 12 ++ 7 files changed, 272 insertions(+), 119 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 5b1d5b9f1..b61d05982 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -849,38 +849,34 @@ class Pregel( # channels are guaranteed to be immutable for the duration of the step, # with channel updates applied only at the transition between steps while loop.tick( - interrupt_before=interrupt_before, interrupt_after=interrupt_after + output_keys=output_keys, + interrupt_before=interrupt_before, + interrupt_after=interrupt_after, ): - # print debug output + # debug flag if self.debug: print_step_checkpoint( loop.checkpoint_metadata, loop.channels, self.stream_channels_list, ) - # emit debug output - if self.checkpointer and "debug" in stream_modes: - yield from _with_mode( - "debug", - isinstance(stream_mode, list), - map_debug_checkpoint( - loop.checkpoint_metadata["step"], - loop.config, - loop.channels, - self.stream_channels_asis, - loop.checkpoint_metadata, - ), - ) - + # emit output + while loop.stream: + mode, payload = loop.stream.popleft() + if mode in stream_modes: + if isinstance(stream_mode, list): + yield (mode, payload) + else: + yield payload + # debug flag if debug: - print_step_tasks(loop.checkpoint_metadata["step"], loop.tasks) + print_step_tasks(loop.step, loop.tasks) + # TODO move to tick() ? if "debug" in stream_modes: yield from _with_mode( "debug", isinstance(stream_mode, list), - map_debug_tasks( - loop.checkpoint_metadata["step"], loop.tasks - ), + map_debug_tasks(loop.step, loop.tasks), ) # execute tasks, and wait for one to fail or all to finish. @@ -931,7 +927,7 @@ class Pregel( "debug", isinstance(stream_mode, list), map_debug_task_results( - loop.checkpoint_metadata["step"], + loop.step, [task], self.stream_channels_list, ), @@ -941,38 +937,31 @@ class Pregel( del fut, task # panic on failure or timeout - _panic_or_proceed(done, inflight, loop.checkpoint_metadata["step"]) + _panic_or_proceed(done, inflight, loop.step) # don't keep futures around in memory longer than needed del done, inflight, futures - - # combine pending writes from all tasks - pending_writes = deque[tuple[str, Any]]() - for task in loop.tasks: - pending_writes.extend(task.writes) - + # debug flag if debug: print_step_writes( - loop.checkpoint_metadata["step"], - pending_writes, + loop.step, + [w for t in loop.tasks for w in t.writes], self.stream_channels_list, ) - - # yield values output - if "values" in stream_modes: - yield from _with_mode( - "values", - isinstance(stream_mode, list), - map_output_values( - output_keys, pending_writes, loop.channels - ), - ) + # emit output + while loop.stream: + mode, payload = loop.stream.popleft() + if mode in stream_modes: + if isinstance(stream_mode, list): + yield (mode, payload) + else: + yield payload + # handle exit if loop.status == "out_of_steps": raise GraphRecursionError( - f"Recursion limit of {config['recursion_limit']} reached" + f"Recursion limit of {config['recursion_limit']} reached " "without hitting a stop condition. You can increase the " "limit by setting the `recursion_limit` config key." ) - # set final channel values as run output run_manager.on_chain_end(read_channels(loop.channels, output_keys)) except BaseException as e: @@ -1142,6 +1131,7 @@ class Pregel( ) def put_checkpoint(metadata: CheckpointMetadata) -> Iterator[Any]: + print(metadata) nonlocal checkpoint, checkpoint_config, channels if self.checkpointer is None: @@ -1268,12 +1258,7 @@ class Pregel( break # before execution, check if we should interrupt - if should_interrupt( - checkpoint, - interrupt_before, - self.stream_channels_list, - next_tasks, - ): + if should_interrupt(checkpoint, interrupt_before, next_tasks): break else: checkpoint = next_checkpoint @@ -1405,12 +1390,7 @@ class Pregel( yield chunk # after execution, check if we should interrupt - if should_interrupt( - checkpoint, - interrupt_after, - self.stream_channels_list, - next_tasks, - ): + if should_interrupt(checkpoint, interrupt_after, next_tasks): break else: raise GraphRecursionError( diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 256e62fd8..46950b9f8 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -44,7 +44,6 @@ from langgraph.pregel.types import All, PregelExecutableTask, PregelTaskDescript def should_interrupt( checkpoint: Checkpoint, interrupt_nodes: Union[All, Sequence[str]], - snapshot_channels: Sequence[str], tasks: list[PregelExecutableTask], ) -> bool: version_type = type(next(iter(checkpoint["channel_versions"].values()), None)) diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 5a317bb26..33b7f440e 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -70,6 +70,8 @@ def map_debug_tasks( if config is not None and TAG_HIDDEN in config.get("tags", []): continue + print("map_debug_tasks", json.dumps((name, step, config["metadata"]))) + yield { "type": "task", "timestamp": ts, @@ -95,6 +97,8 @@ def map_debug_task_results( if config is not None and TAG_HIDDEN in config.get("tags", []): continue + print("map_debug_tasks_r", json.dumps((name, step, config["metadata"]))) + yield { "type": "task_result", "timestamp": ts, diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index f217e201a..d618ffa24 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -4,13 +4,17 @@ from types import TracebackType from typing import ( TYPE_CHECKING, Any, + Callable, ContextManager, List, Literal, Mapping, Optional, Sequence, + Tuple, Type, + TypeVar, + Union, ) from langchain_core.runnables import RunnableConfig @@ -35,31 +39,39 @@ from langgraph.pregel.algo import ( prepare_next_tasks, should_interrupt, ) +from langgraph.pregel.debug import map_debug_checkpoint from langgraph.pregel.executor import BackgroundExecutor, Submit -from langgraph.pregel.io import map_input +from langgraph.pregel.io import map_input, map_output_updates, map_output_values, single from langgraph.pregel.types import PregelExecutableTask if TYPE_CHECKING: from langgraph.pregel import Pregel +V = TypeVar("V") INPUT_DONE = object() class PregelLoop(ContextManager): + input: Optional[Any] config: RunnableConfig - checkpoint: Checkpoint - checkpoint_metadata: CheckpointMetadata - checkpoint_pending_writes: Optional[List[PendingWrite]] + checkpointer: Optional[BaseCheckpointSaver] + get_next_version: Callable[[Optional[V]], V] + graph: "Pregel" submit: Submit channels: Mapping[str, BaseChannel] managed: ManagedValueMapping + checkpoint: Checkpoint + checkpoint_config: RunnableConfig + checkpoint_metadata: CheckpointMetadata + checkpoint_pending_writes: Optional[List[PendingWrite]] status: Literal[ "pending", "done", "interrupt_before", "interrupt_after", "out_of_steps" ] tasks: Sequence[PregelExecutableTask] + stream: deque[Tuple[str, Any]] def __init__( self, @@ -69,6 +81,7 @@ class PregelLoop(ContextManager): checkpointer: Optional[BaseCheckpointSaver], graph: "Pregel", ) -> None: + self.stream = deque() self.stack = ExitStack() self.input = input self.config = config @@ -84,7 +97,7 @@ class PregelLoop(ContextManager): saved = ( self.checkpointer.get_tuple(self.config) if self.checkpointer else None ) or CheckpointTuple(self.config, empty_checkpoint(), {"step": -2}, None, []) - self.config = { + self.checkpoint_config = { **self.config, **saved.config, "configurable": { @@ -106,6 +119,7 @@ class PregelLoop(ContextManager): ) ) self.status = "pending" + self.step = self.checkpoint_metadata["step"] + 1 return self @@ -121,10 +135,11 @@ class PregelLoop(ContextManager): def tick( self, *, + output_keys: Union[str, Sequence[str]] = None, interrupt_after: Optional[Sequence[str]] = None, interrupt_before: Optional[Sequence[str]] = None, ) -> bool: - print("ticking", self.status, self.checkpoint_metadata["step"]) + print("tick", self.status, self.step) if self.status != "pending": raise RuntimeError(f"Cannot tick when status is {self.status}") if self.input is not INPUT_DONE: @@ -132,42 +147,57 @@ class PregelLoop(ContextManager): elif len({tid for tid, _, _ in self.checkpoint_pending_writes}) == len( self.tasks ): + # assign writes to tasks, apply them in order + grouped: dict[str, list[tuple[str, Any]]] = {} + for tid, k, v in self.checkpoint_pending_writes: + grouped.setdefault(tid, []).append((k, v)) + writes = [(k, v) for t in self.tasks for k, v in grouped.get(t.id, [])] # all tasks have finished apply_writes( self.checkpoint, self.channels, - ((k, v) for _, k, v in self.checkpoint_pending_writes), + writes, self.get_next_version, ) + # produce values output + self.stream.extend( + ("values", v) + for v in map_output_values(output_keys, writes, self.channels) + ) # clear pending writes self.checkpoint_pending_writes.clear() # save checkpoint - self.put_checkpoint({"source": "loop", "writes": None}) # TODO + self.put_checkpoint( + { + "source": "loop", + "writes": single( + map_output_updates(output_keys, self.tasks) + if self.graph.stream_mode == "updates" + else map_output_values(output_keys, writes, self.channels) + ), + } + ) # after execution, check if we should interrupt - if should_interrupt( - self.checkpoint, - interrupt_after, - self.graph.stream_channels_list, - self.tasks, - ): + if should_interrupt(self.checkpoint, interrupt_after, self.tasks): self.status = "interrupt_after" return False else: return False # check if iteration limit is reached - if self.checkpoint_metadata["step"] >= self.config["recursion_limit"]: + if self.step > self.config["recursion_limit"]: self.status = "out_of_steps" return False # prepare next tasks + prev_checkpoint = self.checkpoint self.checkpoint, self.tasks = prepare_next_tasks( self.checkpoint, self.graph.nodes, self.channels, self.managed, self.config, - self.checkpoint_metadata["step"], + self.step, for_execution=True, get_next_version=self.get_next_version, ) @@ -180,21 +210,14 @@ class PregelLoop(ContextManager): # TODO how to make this work for both # - online case: we should schedule remaining tasks # - offline case: we should just bail, as other tasks were scheduled before - # assign pending writes to tasks - # if self.checkpoint_pending_writes: - # # if there are pending writes from a previous loop, apply them - # for tid, k, v in self.checkpoint_pending_writes: - # if task := next((t for t in self.tasks if t.id == tid), None): - # task.writes.append((k, v)) - # if + # if there are pending writes from a previous loop, apply them + if self.checkpoint_pending_writes: + for tid, k, v in self.checkpoint_pending_writes: + if task := next((t for t in self.tasks if t.id == tid), None): + task.writes.append((k, v)) # before execution, check if we should interrupt - if should_interrupt( - self.checkpoint, - interrupt_before, - self.graph.stream_channels_list, - self.tasks, - ): + if should_interrupt(prev_checkpoint, interrupt_before, self.tasks): self.status = "interrupt_before" return False @@ -210,9 +233,10 @@ class PregelLoop(ContextManager): self.channels, self.managed, self.config, - -1, + self.step, for_execution=True, get_next_version=self.get_next_version, + # TODO missing run_manager ) # apply input writes apply_writes( @@ -239,9 +263,9 @@ class PregelLoop(ContextManager): self.submit( self.checkpointer.put_writes, { - **self.config, + **self.checkpoint_config, "configurable": { - **self.config["configurable"], + **self.checkpoint_config["configurable"], "thread_ts": self.checkpoint["id"], }, }, @@ -253,30 +277,39 @@ class PregelLoop(ContextManager): self, metadata: CheckpointMetadata, ) -> None: - # increment step - self.checkpoint_metadata = { - **metadata, - "step": self.checkpoint_metadata["step"] + 1, - } + # assign step + metadata["step"] = self.step # bail if no checkpointer - if self.checkpointer is None: - return - # create new checkpoint - self.checkpoint = create_checkpoint( - self.checkpoint, self.channels, self.checkpoint_metadata["step"] - ) - # save it, without blocking - self.submit( - self.checkpointer.put, - self.config, - copy_checkpoint(self.checkpoint), - self.checkpoint_metadata, - ) - # update checkpoint config - self.config = { - **self.config, - "configurable": { - **self.config["configurable"], - "thread_ts": self.checkpoint["id"], - }, - } + if self.checkpointer is not None: + # create new checkpoint + self.checkpoint_metadata = metadata + self.checkpoint = create_checkpoint( + self.checkpoint, self.channels, self.step + ) + # save it, without blocking + self.submit( + self.checkpointer.put, + self.checkpoint_config, + copy_checkpoint(self.checkpoint), + self.checkpoint_metadata, + ) + self.checkpoint_config = { + **self.checkpoint_config, + "configurable": { + **self.checkpoint_config["configurable"], + "thread_ts": self.checkpoint["id"], + }, + } + # produce debug output + self.stream.extend( + ("debug", v) + for v in map_debug_checkpoint( + self.step, + self.checkpoint_config, + self.channels, + self.graph.stream_channels_asis, + self.checkpoint_metadata, + ) + ) + # increment step + self.step += 1 diff --git a/libs/langgraph/tests/test_prebuilt.py b/libs/langgraph/tests/test_prebuilt.py index 6053a4e5b..7b56eba46 100644 --- a/libs/langgraph/tests/test_prebuilt.py +++ b/libs/langgraph/tests/test_prebuilt.py @@ -1,5 +1,8 @@ +from collections import defaultdict from typing import Any, Callable, Dict, List, Optional, Sequence, Type, Union +from langgraph.checkpoint.base import BaseCheckpointSaver +from langgraph.checkpoint.sqlite import SqliteSaver import pytest from langchain_core.callbacks import ( CallbackManagerForLLMRun, @@ -27,6 +30,8 @@ from langgraph.prebuilt import ( ValidationNode, create_react_agent, ) +from tests.any_str import AnyStr +from tests.memory_assert import MemorySaverAssertImmutable class FakeToolCallingModel(BaseChatModel): @@ -56,14 +61,121 @@ class FakeToolCallingModel(BaseChatModel): return self -def test_no_modifier(): +@pytest.mark.parametrize( + "checkpointer", + [ + MemorySaverAssertImmutable(), + None, + ], + ids=[ + "memory", + "none", + ], +) +def test_no_modifier(checkpointer: Optional[BaseCheckpointSaver]): model = FakeToolCallingModel() - agent = create_react_agent(model, []) + agent = create_react_agent(model, [], checkpointer=checkpointer) inputs = [HumanMessage("hi?")] - response = agent.invoke({"messages": inputs}, debug=True) + thread = {"configurable": {"thread_id": "123"}} + response = agent.invoke({"messages": inputs}, thread, debug=True) expected_response = {"messages": inputs + [AIMessage(content="hi?", id="0")]} assert response == expected_response + if checkpointer: + saved = checkpointer.get_tuple(thread) + assert saved is not None + assert saved.checkpoint == { + "v": 1, + "ts": AnyStr(), + "id": AnyStr(), + "channel_values": { + "messages": [ + HumanMessage(content="hi?", id=AnyStr()), + AIMessage(content="hi?", id="0"), + ], + "agent": "agent", + }, + "channel_versions": { + "__start__": 2, + "messages": 3, + "start:agent": 3, + "agent": 3, + }, + "versions_seen": defaultdict( + dict, + { + "__start__": {"__start__": 1}, + "agent": {"start:agent": 2}, + "tools": {}, + }, + ), + "pending_sends": [], + } + assert saved.metadata == { + "source": "loop", + "writes": {"agent": {"messages": [AIMessage(content="hi?", id="0")]}}, + "step": 1, + } + assert saved.pending_writes == [] + + +@pytest.mark.parametrize( + "checkpointer", + [ + MemorySaverAssertImmutable(), + None, + ], + ids=[ + "memory", + "none", + ], +) +async def test_no_modifier_async(checkpointer: Optional[BaseCheckpointSaver]): + model = FakeToolCallingModel() + agent = create_react_agent(model, [], checkpointer=checkpointer) + inputs = [HumanMessage("hi?")] + thread = {"configurable": {"thread_id": "123"}} + response = await agent.ainvoke({"messages": inputs}, thread, debug=True) + expected_response = {"messages": inputs + [AIMessage(content="hi?", id="0")]} + assert response == expected_response + + if checkpointer: + saved = await checkpointer.aget_tuple(thread) + assert saved is not None + assert saved.checkpoint == { + "v": 1, + "ts": AnyStr(), + "id": AnyStr(), + "channel_values": { + "messages": [ + HumanMessage(content="hi?", id=AnyStr()), + AIMessage(content="hi?", id="0"), + ], + "agent": "agent", + }, + "channel_versions": { + "__start__": 2, + "messages": 3, + "start:agent": 3, + "agent": 3, + }, + "versions_seen": defaultdict( + dict, + { + "__start__": {"__start__": 1}, + "agent": {"start:agent": 2}, + "tools": {}, + }, + ), + "pending_sends": [], + } + assert saved.metadata == { + "source": "loop", + "writes": {"agent": {"messages": [AIMessage(content="hi?", id="0")]}}, + "step": 1, + } + assert saved.pending_writes == [] + def test_passing_two_modifiers(): model = FakeToolCallingModel() diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 805794bdf..9b32fc149 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -521,7 +521,7 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert app.invoke(2) == 4 with pytest.raises(GraphRecursionError): - app.invoke(2, {"recursion_limit": 1}) + app.invoke(2, {"recursion_limit": 1}, debug=1) graph = Graph() graph.add_node("add_one", add_one) @@ -533,7 +533,7 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert gapp.invoke(2) == 4 - for step, values in enumerate(gapp.stream(2), start=1): + for step, values in enumerate(gapp.stream(2, debug=1), start=1): if step == 1: assert values == { "add_one": 3, @@ -6166,6 +6166,18 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: "my_key": "value ⛰️", "market": "DE", } + assert [c.metadata for c in tool_two.checkpointer.list(thread1)] == [ + { + "source": "loop", + "step": 0, + "writes": None, + }, + { + "source": "input", + "step": -1, + "writes": {"my_key": "value ⛰️", "market": "DE"}, + }, + ] assert tool_two.get_state(thread1) == StateSnapshot( values={"my_key": "value ⛰️", "market": "DE"}, next=("tool_two_slow",), @@ -6875,7 +6887,8 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> }, ) - assert [c for c in app_w_interrupt.stream(None, config)] == [ + print("yo") + assert [c for c in app_w_interrupt.stream(None, config, debug=1)] == [ {"qa": {"answer": "doc1,doc2,doc3,doc4,doc5"}}, ] diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 1cf5c35d4..0e927e518 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -4735,6 +4735,18 @@ async def test_start_branch_then() -> None: "my_key": "value", "market": "DE", } + assert [c.metadata async for c in tool_two.checkpointer.alist(thread1)] == [ + { + "source": "loop", + "step": 0, + "writes": None, + }, + { + "source": "input", + "step": -1, + "writes": {"my_key": "value", "market": "DE"}, + }, + ] assert await tool_two.aget_state(thread1) == StateSnapshot( values={"my_key": "value", "market": "DE"}, next=("tool_two_slow",),