From fdb273b0f430c47ecaa845297f82b25d613aae71 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 20 Feb 2024 09:52:31 -0800 Subject: [PATCH 1/8] Add get_state() and update_state() methods to get and update checkpoint in between runs --- Makefile | 2 +- langgraph/graph/graph.py | 25 ++- langgraph/graph/state.py | 47 +++++- langgraph/pregel/__init__.py | 95 ++++++++++- pyproject.toml | 6 + tests/test_pregel.py | 310 +++++++++++++++++++++++++++++++++- tests/test_pregel_async.py | 316 ++++++++++++++++++++++++++++++++++- 7 files changed, 786 insertions(+), 15 deletions(-) diff --git a/Makefile b/Makefile index 6a4b86629..c50d63a07 100644 --- a/Makefile +++ b/Makefile @@ -18,7 +18,7 @@ test: poetry run pytest test_watch: - poetry run ptw --snapshot-update --now . -- -vv -x --ff tests + poetry run ptw tests ###################### # LINTING AND FORMATTING diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 4885f3b6c..afbe9e1fd 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -11,9 +11,10 @@ from langchain_core.runnables.base import ( ) from langchain_core.runnables.config import RunnableConfig from langchain_core.runnables.graph import Graph as RunnableGraph +from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.checkpoint import BaseCheckpointSaver -from langgraph.pregel import Channel, Pregel +from langgraph.pregel import Channel, Pregel, StateSnapshot logger = logging.getLogger(__name__) @@ -190,6 +191,11 @@ class Graph: key: (Channel.subscribe_to(f"{key}:inbox") | node | Channel.write_to(key)) for key, node in self.nodes.items() } + node_outboxes = { + # we clear outbox channels after each step + key: EphemeralValue(Any) + for key in self.nodes + } for key in self.nodes: outgoing = outgoing_edges[key] @@ -216,6 +222,7 @@ class Graph: return CompiledGraph( graph=self, nodes=nodes, + channels={**node_outboxes}, input=f"{self.entry_point}:inbox" if self.entry_point else START, output=END, hidden=[f"{node}:inbox" for node in self.nodes], @@ -272,3 +279,19 @@ class CompiledGraph(Pregel): graph.add_edge(graph.nodes[START], graph.nodes[self.graph.entry_point]) return graph + + def get_state(self, config: RunnableConfig) -> StateSnapshot: + snapshot = super().get_state(config) + + return StateSnapshot( + values={k: v for k, v in snapshot.values.items() if k in self.graph.nodes}, + next=snapshot.next, + ) + + async def aget_state(self, config: RunnableConfig) -> StateSnapshot: + snapshot = await super().aget_state(config) + + return StateSnapshot( + values={k: v for k, v in snapshot.values.items() if k in self.graph.nodes}, + next=snapshot.next, + ) diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index c448b7628..bfc459a1c 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -1,10 +1,11 @@ from collections import defaultdict from functools import partial from inspect import signature -from typing import Any, Optional, Sequence, Type +from typing import Any, Optional, Sequence, Type, Union from langchain_core.runnables import RunnableLambda from langchain_core.runnables.base import RunnableLike +from langchain_core.runnables.config import RunnableConfig from langgraph.channels.any_value import AnyValue from langgraph.channels.base import BaseChannel, InvalidUpdateError @@ -13,7 +14,7 @@ from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.checkpoint import BaseCheckpointSaver from langgraph.graph.graph import END, START, CompiledGraph, Graph -from langgraph.pregel import Channel +from langgraph.pregel import Channel, StateSnapshot from langgraph.pregel.read import ChannelInvoke from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry @@ -134,7 +135,7 @@ class StateGraph(Graph): else: raise ValueError("No entry point set") - return CompiledGraph( + return CompiledStateGraph( graph=self, nodes=nodes, channels={ @@ -204,3 +205,43 @@ def _is_field_binop(typ: Type[Any]) -> Optional[BinaryOperatorAggregate]: ): return BinaryOperatorAggregate(typ, meta[0]) return None + + +class CompiledStateGraph(CompiledGraph): + graph: StateGraph + + def get_state(self, config: RunnableConfig) -> StateSnapshot: + snapshot = super(CompiledGraph, self).get_state(config) + + return StateSnapshot( + values=snapshot.values.get("__root__") + if "__root__" in self.graph.channels + else {k: v for k, v in snapshot.values.items() if k in self.graph.channels}, + next=snapshot.next, + ) + + async def aget_state(self, config: RunnableConfig) -> StateSnapshot: + snapshot = await super(CompiledGraph, self).aget_state(config) + + return StateSnapshot( + values=snapshot.values.get("__root__") + if "__root__" in self.graph.channels + else {k: v for k, v in snapshot.values.items() if k in self.graph.channels}, + next=snapshot.next, + ) + + def update_state( + self, config: RunnableConfig, values: Union[Any, dict[str, Any]] + ) -> None: + return super(CompiledGraph, self).update_state( + config, + {"__root__": values} if "__root__" in self.graph.channels else values, + ) + + async def aupdate_state( + self, config: RunnableConfig, values: Union[Any, dict[str, Any]] + ) -> None: + return await super(CompiledGraph, self).aupdate_state( + config, + {"__root__": values} if "__root__" in self.graph.channels else values, + ) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index d512bd3dd..4f443413d 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -11,6 +11,7 @@ from typing import ( Callable, Iterator, Mapping, + NamedTuple, Optional, Sequence, Type, @@ -40,6 +41,7 @@ from langchain_core.runnables.utils import ( get_unique_config_specs, ) from langchain_core.tracers.log_stream import LogStreamCallbackHandler +from langgraph.channels.any_value import AnyValue from langgraph.channels.base import ( AsyncChannelsManager, @@ -49,6 +51,7 @@ from langgraph.channels.base import ( InvalidUpdateError, create_checkpoint, ) +from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.checkpoint.base import ( BaseCheckpointSaver, @@ -158,6 +161,13 @@ class Channel: ) +class StateSnapshot(NamedTuple): + values: dict[str, Any] + """Current values of channels""" + next: tuple[str] + """Nodes to execute in the next step, if any""" + + class Pregel( RunnableSerializable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]] ): @@ -247,6 +257,72 @@ class Pregel( **{k: (self.channels[k].ValueType, None) for k in self.output}, ) + def get_state(self, config: RunnableConfig) -> StateSnapshot: + if not self.checkpointer: + raise ValueError("No checkpointer set") + + checkpoint = self.checkpointer.get(config) + checkpoint = checkpoint or empty_checkpoint() + with ChannelsManager(self.channels, checkpoint) as channels: + next_tasks = _prepare_next_tasks( + checkpoint, self.nodes, channels, update_seen=False + ) + return StateSnapshot( + { + k: _read_channel(channels, k) + for k in channels + if k not in [k.value for k in ReservedChannels] + }, + tuple(name for _, _, name in next_tasks), + ) + + async def aget_state(self, config: RunnableConfig) -> StateSnapshot: + if not self.checkpointer: + raise ValueError("No checkpointer set") + + checkpoint = await self.checkpointer.aget(config) + checkpoint = checkpoint or empty_checkpoint() + async with AsyncChannelsManager(self.channels, checkpoint) as channels: + next_tasks = _prepare_next_tasks( + checkpoint, self.nodes, channels, update_seen=False + ) + return StateSnapshot( + { + k: _read_channel(channels, k) + for k in channels + if k not in [k.value for k in ReservedChannels] + }, + tuple(name for _, _, name in next_tasks), + ) + + def update_state(self, config: RunnableConfig, values: dict[str, Any]) -> None: + if not self.checkpointer: + raise ValueError("No checkpointer set") + + checkpoint = self.checkpointer.get(config) + checkpoint = checkpoint or empty_checkpoint() + with ChannelsManager(self.channels, checkpoint) as channels: + for k, v in values.items(): + channels[k].update([v]) + checkpoint["channel_versions"][k] += 1 + self.checkpointer.put(config, create_checkpoint(checkpoint, channels)) + + async def aupdate_state( + self, config: RunnableConfig, values: dict[str, Any] + ) -> None: + if not self.checkpointer: + raise ValueError("No checkpointer set") + + checkpoint = await self.checkpointer.aget(config) + checkpoint = checkpoint or empty_checkpoint() + async with AsyncChannelsManager(self.channels, checkpoint) as channels: + for k, v in values.items(): + channels[k].update([v]) + checkpoint["channel_versions"][k] += 1 + await self.checkpointer.aput( + config, create_checkpoint(checkpoint, channels) + ) + def _transform( self, input: Iterator[Union[dict[str, Any], Any]], @@ -768,7 +844,7 @@ def _apply_writes_from_view( if value == _read_channel(channels, chan): continue - assert isinstance(channels[chan], LastValue), ( + assert isinstance(channels[chan], (LastValue, EphemeralValue, AnyValue)), ( f"Can't modify channel {chan} of type " f"{channels[chan].__class__.__name__}" ) @@ -780,6 +856,7 @@ def _prepare_next_tasks( checkpoint: Checkpoint, processes: Mapping[str, Union[ChannelInvoke, ChannelBatch]], channels: Mapping[str, BaseChannel], + update_seen: bool = True, ) -> list[tuple[Runnable, Any, str]]: tasks: list[tuple[Runnable, Any, str]] = [] # Check if any processes should be run in next step @@ -814,12 +891,13 @@ def _prepare_next_tasks( val = val[None] # update seen versions - seen.update( - { - chan: checkpoint["channel_versions"][chan] - for chan in proc.triggers - } - ) + if update_seen: + seen.update( + { + chan: checkpoint["channel_versions"][chan] + for chan in proc.triggers + } + ) # skip if condition is not met if proc.when is None or proc.when(val): @@ -836,7 +914,8 @@ def _prepare_next_tasks( val = [{proc.key: v} for v in val] tasks.append((proc, val, name)) - seen[proc.channel] = checkpoint["channel_versions"][proc.channel] + if update_seen: + seen[proc.channel] = checkpoint["channel_versions"][proc.channel] return tasks diff --git a/pyproject.toml b/pyproject.toml index 723c1ac21..9e2d94a0b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,12 @@ exclude = ["notebooks", "examples", "example_data"] [tool.coverage.run] omit = ["tests/*"] +[tool.pytest-watcher] +now = true +delay = 0.1 +runner_args = ["-x", "--ff", "-vv", "--snapshot-update"] +patterns = ["*.py"] + [build-system] requires = ["poetry-core>=1.0.0"] build-backend = "poetry.core.masonry.api" diff --git a/tests/test_pregel.py b/tests/test_pregel.py index d1cb08a02..6ae5f5ee9 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -26,7 +26,7 @@ from langgraph.prebuilt.chat_agent_executor import ( create_tool_calling_executor, ) from langgraph.prebuilt.tool_executor import ToolExecutor -from langgraph.pregel import Channel, GraphRecursionError, Pregel +from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.reserved import ReservedChannels @@ -282,6 +282,22 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: assert app.invoke(3, {"configurable": {"thread_id": 1}}) is None assert app.invoke(None, {"configurable": {"thread_id": 1}}) == 5 + # start execution again, stopping at inbox + assert app.invoke(20, {"configurable": {"thread_id": 2}}) is None + + # inbox == 21 + snapshot = app.get_state({"configurable": {"thread_id": 2}}) + assert snapshot.values["inbox"] == 21 + assert snapshot.next == ("two",) + + # update the state, resume + app.update_state({"configurable": {"thread_id": 2}}, {"inbox": 25}) + assert app.invoke(None, {"configurable": {"thread_id": 2}}) == 26 + + # no pending tasks + snapshot = app.get_state({"configurable": {"thread_id": 2}}) + assert snapshot.next == () + def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) @@ -761,6 +777,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: ), } + # deepcopy because the nodes mutate the data assert [deepcopy(c) for c in app.stream({"input": "what is weather in sf"})] == [ { "agent": { @@ -882,6 +899,151 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: }, ] + # test state get/update methods + + app_w_interrupt = workflow.compile( + checkpointer=MemorySaver(), interrupt_after=["agent"] + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config) + ] == [ + { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + } + } + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + }, + "tools": None, + }, + next=("agent:edges",), + ) + + app_w_interrupt.update_state( + config, + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + }, + }, + ) + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + }, + "tools": None, + }, + next=("agent:edges",), + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "tools": { + "input": "what is weather in sf", + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ) + ], + } + }, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ) + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + ] + + app_w_interrupt.update_state( + config, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ) + ], + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + } + }, + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "__end__": { + "input": "what is weather in sf", + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ) + ], + } + } + ] + def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None: from langchain.llms.fake import FakeStreamingListLLM @@ -1073,6 +1235,116 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None: }, ] + app_w_interrupt = workflow.compile( + checkpointer=MemorySaver(), interrupt_after=["agent"] + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config) + ] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + } + } + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + "intermediate_steps": [], + }, + next=("agent:edges",), + ) + + app_w_interrupt.update_state( + config, + { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ) + }, + ) + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "intermediate_steps": [], + }, + next=("agent:edges",), + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "tools": { + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ) + ], + } + }, + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + ] + + app_w_interrupt.update_state( + config, + { + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ) + }, + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "__end__": { + "input": "what is weather in sf", + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ) + ], + } + } + ] + def test_conditional_entrypoint_graph(snapshot: SnapshotAssertion) -> None: def left(data: str) -> str: @@ -1749,6 +2021,42 @@ def test_message_graph(snapshot: SnapshotAssertion) -> None: }, ] + app_w_interrupt = workflow.compile( + checkpointer=MemorySaver(), interrupt_after=["agent"] + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + for c in app_w_interrupt.stream( + HumanMessage(content="what is weather in sf"), config + ) + ] == [ + { + "agent": AIMessage( + content="", + additional_kwargs={ + "function_call": {"name": "search_api", "arguments": '"query"'} + }, + ) + } + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values=[ + HumanMessage(content="what is weather in sf"), + AIMessage( + content="", + additional_kwargs={ + "function_call": {"name": "search_api", "arguments": '"query"'} + }, + ), + ], + next=("agent:edges",), + ) + + # TODO use update_state once we have message ids + def test_in_one_fan_out_out_one_graph_state() -> None: def sorted_add(x: list[str], y: list[str]) -> list[str]: diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index d7b762c34..6db086683 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -31,7 +31,7 @@ from langgraph.prebuilt.chat_agent_executor import ( create_tool_calling_executor, ) from langgraph.prebuilt.tool_executor import ToolExecutor -from langgraph.pregel import Channel, GraphRecursionError, Pregel +from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.reserved import ReservedChannels @@ -289,6 +289,22 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N assert await app.ainvoke(3, {"configurable": {"thread_id": 1}}) is None assert await app.ainvoke(None, {"configurable": {"thread_id": 1}}) == 5 + # start execution again, stopping at inbox + assert await app.ainvoke(20, {"configurable": {"thread_id": 2}}) is None + + # inbox == 21 + snapshot = await app.aget_state({"configurable": {"thread_id": 2}}) + assert snapshot.values["inbox"] == 21 + assert snapshot.next == ("two",) + + # update the state, resume + await app.aupdate_state({"configurable": {"thread_id": 2}}, {"inbox": 25}) + assert await app.ainvoke(None, {"configurable": {"thread_id": 2}}) == 26 + + # no pending tasks + snapshot = await app.aget_state({"configurable": {"thread_id": 2}}) + assert snapshot.next == () + async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) @@ -798,6 +814,7 @@ async def test_conditional_graph() -> None: ), } + # deepcopy because the nodes mutate the data assert [ deepcopy(c) async for c in app.astream({"input": "what is weather in sf"}) ] == [ @@ -927,6 +944,154 @@ async def test_conditional_graph() -> None: # Check that agent (one of the nodes) has its output streamed to the logs assert "/logs/agent/streamed_output/-" in patch_paths + # test state get/update methods + + app_w_interrupt = workflow.compile( + checkpointer=MemorySaver(), interrupt_after=["agent"] + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + {"input": "what is weather in sf"}, config + ) + ] == [ + { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + } + } + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + }, + "tools": None, + }, + next=("agent:edges",), + ) + + await app_w_interrupt.aupdate_state( + config, + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + }, + }, + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + }, + "tools": None, + }, + next=("agent:edges",), + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": { + "input": "what is weather in sf", + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ) + ], + } + }, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ) + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + ] + + await app_w_interrupt.aupdate_state( + config, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ) + ], + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + } + }, + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "__end__": { + "input": "what is weather in sf", + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ) + ], + } + } + ] + async def test_conditional_graph_state() -> None: from langchain.llms.fake import FakeStreamingListLLM @@ -1113,6 +1278,119 @@ async def test_conditional_graph_state() -> None: }, ] + app_w_interrupt = workflow.compile( + checkpointer=MemorySaver(), interrupt_after=["agent"] + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + {"input": "what is weather in sf"}, config + ) + ] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + } + } + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + "intermediate_steps": [], + }, + next=("agent:edges",), + ) + + await app_w_interrupt.aupdate_state( + config, + { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ) + }, + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "intermediate_steps": [], + }, + next=("agent:edges",), + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": { + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ) + ], + } + }, + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + ] + + await app_w_interrupt.aupdate_state( + config, + { + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ) + }, + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "__end__": { + "input": "what is weather in sf", + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ) + ], + } + } + ] + async def test_conditional_entrypoint_graph() -> None: async def left(data: str) -> str: @@ -1772,6 +2050,42 @@ async def test_message_graph() -> None: }, ] + app_w_interrupt = workflow.compile( + checkpointer=MemorySaver(), interrupt_after=["agent"] + ) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c + async for c in app_w_interrupt.astream( + HumanMessage(content="what is weather in sf"), config + ) + ] == [ + { + "agent": AIMessage( + content="", + additional_kwargs={ + "function_call": {"name": "search_api", "arguments": '"query"'} + }, + ) + } + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values=[ + HumanMessage(content="what is weather in sf"), + AIMessage( + content="", + additional_kwargs={ + "function_call": {"name": "search_api", "arguments": '"query"'} + }, + ), + ], + next=("agent:edges",), + ) + + # TODO use update_state once we have message ids + async def test_in_one_fan_out_out_one_graph_state() -> None: def sorted_add(x: list[str], y: list[str]) -> list[str]: From 38cd934f380157248230940fe78a778e8f172e9b Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 20 Feb 2024 17:23:54 -0800 Subject: [PATCH 2/8] Lint --- langgraph/graph/graph.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index afbe9e1fd..56e1afa33 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -11,8 +11,8 @@ from langchain_core.runnables.base import ( ) from langchain_core.runnables.config import RunnableConfig from langchain_core.runnables.graph import Graph as RunnableGraph -from langgraph.channels.ephemeral_value import EphemeralValue +from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.checkpoint import BaseCheckpointSaver from langgraph.pregel import Channel, Pregel, StateSnapshot From 82c704e57e21812271a96b7d0f4b9db814c119f7 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sat, 24 Feb 2024 17:00:45 -0800 Subject: [PATCH 3/8] Lint --- langgraph/pregel/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 4f443413d..fa061d619 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -41,8 +41,8 @@ from langchain_core.runnables.utils import ( get_unique_config_specs, ) from langchain_core.tracers.log_stream import LogStreamCallbackHandler -from langgraph.channels.any_value import AnyValue +from langgraph.channels.any_value import AnyValue from langgraph.channels.base import ( AsyncChannelsManager, BaseChannel, From 5c2afc94180b36f521429c073264eb8bf32dab1d Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sat, 24 Feb 2024 19:43:30 -0800 Subject: [PATCH 4/8] Fix behaviour of interrupt_before - checkpoints are now copied before being mutated - compiled graph and compiled state graph no longer need to override get_state/update_state - pregel class now natively supports interrupt before/after --- Makefile | 2 +- langgraph/channels/base.py | 16 +- langgraph/checkpoint/base.py | 11 ++ langgraph/constants.py | 1 + langgraph/graph/graph.py | 25 +-- langgraph/graph/state.py | 45 +----- langgraph/pregel/__init__.py | 296 ++++++++++++++++++++++++++--------- langgraph/pregel/validate.py | 11 +- poetry.lock | 11 +- pyproject.toml | 2 +- tests/memory_assert.py | 19 +++ tests/test_pregel.py | 281 ++++++++++++++++++++++++++++++++- tests/test_pregel_async.py | 285 ++++++++++++++++++++++++++++++++- 13 files changed, 832 insertions(+), 173 deletions(-) create mode 100644 tests/memory_assert.py diff --git a/Makefile b/Makefile index c50d63a07..418703a74 100644 --- a/Makefile +++ b/Makefile @@ -18,7 +18,7 @@ test: poetry run pytest test_watch: - poetry run ptw tests + poetry run ptw . ###################### # LINTING AND FORMATTING diff --git a/langgraph/channels/base.py b/langgraph/channels/base.py index b30bee9b2..71680c8ea 100644 --- a/langgraph/channels/base.py +++ b/langgraph/channels/base.py @@ -116,16 +116,16 @@ def create_checkpoint( checkpoint: Checkpoint, channels: Mapping[str, BaseChannel] ) -> Checkpoint: """Create a checkpoint for the given channels.""" - checkpoint = Checkpoint( + values: dict[str, Any] = {} + for k, v in channels.items(): + try: + values[k] = v.checkpoint() + except EmptyChannelError: + pass + return Checkpoint( v=1, ts=datetime.now(timezone.utc).isoformat(), - channel_values=checkpoint["channel_values"], + channel_values=values, channel_versions=checkpoint["channel_versions"], versions_seen=checkpoint["versions_seen"], ) - for k, v in channels.items(): - try: - checkpoint["channel_values"][k] = v.checkpoint() - except EmptyChannelError: - pass - return checkpoint diff --git a/langgraph/checkpoint/base.py b/langgraph/checkpoint/base.py index 699dc9dd8..99b0e294a 100644 --- a/langgraph/checkpoint/base.py +++ b/langgraph/checkpoint/base.py @@ -1,6 +1,7 @@ import asyncio from abc import ABC, abstractmethod from collections import defaultdict +from copy import deepcopy from datetime import datetime, timezone from typing import Any, Optional, TypedDict @@ -33,6 +34,16 @@ def empty_checkpoint() -> Checkpoint: ) +def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint: + return Checkpoint( + v=checkpoint["v"], + ts=checkpoint["ts"], + channel_values=checkpoint["channel_values"].copy(), + channel_versions=checkpoint["channel_versions"].copy(), + versions_seen=deepcopy(checkpoint["versions_seen"]), + ) + + class CheckpointAt(StrEnum): END_OF_STEP = "end_of_step" END_OF_RUN = "end_of_run" diff --git a/langgraph/constants.py b/langgraph/constants.py index 6e5f9c37f..4bf8b1335 100644 --- a/langgraph/constants.py +++ b/langgraph/constants.py @@ -1,2 +1,3 @@ CONFIG_KEY_SEND = "__pregel_send" CONFIG_KEY_READ = "__pregel_read" +INTERRUPT = "__interrupt__" diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 56e1afa33..53331d2af 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -178,6 +178,7 @@ class Graph: checkpointer: Optional[BaseCheckpointSaver] = None, interrupt_before: Optional[Sequence[str]] = None, interrupt_after: Optional[Sequence[str]] = None, + debug: bool = False, ) -> "CompiledGraph": interrupt_before = interrupt_before or [] interrupt_after = interrupt_after or [] @@ -226,11 +227,11 @@ class Graph: input=f"{self.entry_point}:inbox" if self.entry_point else START, output=END, hidden=[f"{node}:inbox" for node in self.nodes], + snapshot_channels=list(self.nodes), checkpointer=checkpointer, - interrupt=( - [f"{node}:inbox" for node in interrupt_before] - + [node for node in interrupt_after] - ), + interrupt_before_nodes=[f"{node}:inbox" for node in interrupt_before], + interrupt_after_nodes=interrupt_after, + debug=debug, ) @@ -279,19 +280,3 @@ class CompiledGraph(Pregel): graph.add_edge(graph.nodes[START], graph.nodes[self.graph.entry_point]) return graph - - def get_state(self, config: RunnableConfig) -> StateSnapshot: - snapshot = super().get_state(config) - - return StateSnapshot( - values={k: v for k, v in snapshot.values.items() if k in self.graph.nodes}, - next=snapshot.next, - ) - - async def aget_state(self, config: RunnableConfig) -> StateSnapshot: - snapshot = await super().aget_state(config) - - return StateSnapshot( - values={k: v for k, v in snapshot.values.items() if k in self.graph.nodes}, - next=snapshot.next, - ) diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index bfc459a1c..b1f048074 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -40,6 +40,7 @@ class StateGraph(Graph): checkpointer: Optional[BaseCheckpointSaver] = None, interrupt_before: Optional[Sequence[str]] = None, interrupt_after: Optional[Sequence[str]] = None, + debug: bool = False, ) -> CompiledGraph: interrupt_before = interrupt_before or [] interrupt_after = interrupt_after or [] @@ -147,11 +148,11 @@ class StateGraph(Graph): input=f"{START}:inbox", output=END, hidden=[f"{node}:inbox" for node in self.nodes] + [START] + state_keys, + snapshot_channels=state_keys_read, checkpointer=checkpointer, - interrupt=( - [f"{node}:inbox" for node in interrupt_before] - + [node for node in interrupt_after] - ), + interrupt_before_nodes=[f"{node}:inbox" for node in interrupt_before], + interrupt_after_nodes=interrupt_after, + debug=debug, ) @@ -209,39 +210,3 @@ def _is_field_binop(typ: Type[Any]) -> Optional[BinaryOperatorAggregate]: class CompiledStateGraph(CompiledGraph): graph: StateGraph - - def get_state(self, config: RunnableConfig) -> StateSnapshot: - snapshot = super(CompiledGraph, self).get_state(config) - - return StateSnapshot( - values=snapshot.values.get("__root__") - if "__root__" in self.graph.channels - else {k: v for k, v in snapshot.values.items() if k in self.graph.channels}, - next=snapshot.next, - ) - - async def aget_state(self, config: RunnableConfig) -> StateSnapshot: - snapshot = await super(CompiledGraph, self).aget_state(config) - - return StateSnapshot( - values=snapshot.values.get("__root__") - if "__root__" in self.graph.channels - else {k: v for k, v in snapshot.values.items() if k in self.graph.channels}, - next=snapshot.next, - ) - - def update_state( - self, config: RunnableConfig, values: Union[Any, dict[str, Any]] - ) -> None: - return super(CompiledGraph, self).update_state( - config, - {"__root__": values} if "__root__" in self.graph.channels else values, - ) - - async def aupdate_state( - self, config: RunnableConfig, values: Union[Any, dict[str, Any]] - ) -> None: - return await super(CompiledGraph, self).aupdate_state( - config, - {"__root__": values} if "__root__" in self.graph.channels else values, - ) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index fa061d619..76df9f591 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -57,9 +57,10 @@ from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, CheckpointAt, + copy_checkpoint, empty_checkpoint, ) -from langgraph.constants import CONFIG_KEY_READ, CONFIG_KEY_SEND +from langgraph.constants import CONFIG_KEY_READ, CONFIG_KEY_SEND, INTERRUPT from langgraph.pregel.debug import print_checkpoint, print_step_start from langgraph.pregel.io import map_input, map_output from langgraph.pregel.log import logger @@ -162,7 +163,7 @@ class Channel: class StateSnapshot(NamedTuple): - values: dict[str, Any] + values: dict[str, Any] | Any """Current values of channels""" next: tuple[str] """Nodes to execute in the next step, if any""" @@ -175,12 +176,19 @@ class Pregel( channels: Mapping[str, BaseChannel] = Field(default_factory=dict) + # TODO Rename to `output_channels` output: Union[str, Sequence[str]] = "output" + # TODO Replace with `stream_channels` hidden: Sequence[str] = Field(default_factory=list) - interrupt: Sequence[str] = Field(default_factory=list) + snapshot_channels: Union[str, Sequence[str]] = Field(default_factory=list) + interrupt_after_nodes: Sequence[str] = Field(default_factory=list) + + interrupt_before_nodes: Sequence[str] = Field(default_factory=list) + + # TODO Rename to `input_channels` input: Union[str, Sequence[str]] = "input" step_timeout: Optional[float] = None @@ -202,8 +210,12 @@ class Pregel( values["input"], values["output"], values["hidden"], - values["interrupt"], + values["interrupt_after_nodes"], + values["interrupt_before_nodes"], ) + if values["interrupt_after_nodes"] or values["interrupt_before_nodes"]: + if not values["checkpointer"]: + raise ValueError("Interrupts require a checkpointer") return values @property @@ -257,6 +269,14 @@ class Pregel( **{k: (self.channels[k].ValueType, None) for k in self.output}, ) + @property + def snapshot_channels_list(self) -> Sequence[str]: + return ( + [self.snapshot_channels] + if isinstance(self.snapshot_channels, str) + else self.snapshot_channels or list(self.channels.keys()) + ) + def get_state(self, config: RunnableConfig) -> StateSnapshot: if not self.checkpointer: raise ValueError("No checkpointer set") @@ -264,15 +284,19 @@ class Pregel( checkpoint = self.checkpointer.get(config) checkpoint = checkpoint or empty_checkpoint() with ChannelsManager(self.channels, checkpoint) as channels: - next_tasks = _prepare_next_tasks( + _, next_tasks = _prepare_next_tasks( checkpoint, self.nodes, channels, update_seen=False ) + values = { + k: _read_channel(channels, k) + for k in channels + if k in self.snapshot_channels_list + and k not in [k.value for k in ReservedChannels] + } return StateSnapshot( - { - k: _read_channel(channels, k) - for k in channels - if k not in [k.value for k in ReservedChannels] - }, + values[self.snapshot_channels] + if isinstance(self.snapshot_channels, str) + else values, tuple(name for _, _, name in next_tasks), ) @@ -283,71 +307,130 @@ class Pregel( checkpoint = await self.checkpointer.aget(config) checkpoint = checkpoint or empty_checkpoint() async with AsyncChannelsManager(self.channels, checkpoint) as channels: - next_tasks = _prepare_next_tasks( + _, next_tasks = _prepare_next_tasks( checkpoint, self.nodes, channels, update_seen=False ) + values = { + k: _read_channel(channels, k) + for k in channels + if k in self.snapshot_channels_list + and k not in [k.value for k in ReservedChannels] + } return StateSnapshot( - { - k: _read_channel(channels, k) - for k in channels - if k not in [k.value for k in ReservedChannels] - }, + values[self.snapshot_channels] + if isinstance(self.snapshot_channels, str) + else values, tuple(name for _, _, name in next_tasks), ) - def update_state(self, config: RunnableConfig, values: dict[str, Any]) -> None: - if not self.checkpointer: - raise ValueError("No checkpointer set") - - checkpoint = self.checkpointer.get(config) - checkpoint = checkpoint or empty_checkpoint() - with ChannelsManager(self.channels, checkpoint) as channels: - for k, v in values.items(): - channels[k].update([v]) - checkpoint["channel_versions"][k] += 1 - self.checkpointer.put(config, create_checkpoint(checkpoint, channels)) - - async def aupdate_state( - self, config: RunnableConfig, values: dict[str, Any] + def update_state( + self, config: RunnableConfig, values: dict[str, Any] | Any ) -> None: if not self.checkpointer: raise ValueError("No checkpointer set") + values = ( + {self.snapshot_channels: values} + if isinstance(self.snapshot_channels, str) + else values + ) + checkpoint = self.checkpointer.get(config) + checkpoint = copy_checkpoint(checkpoint) if checkpoint else empty_checkpoint() + with ChannelsManager(self.channels, checkpoint) as channels: + for k, v in values.items(): + channels[k].update([v]) + checkpoint["channel_versions"][k] += 1 + for k in self.snapshot_channels or self.channels: + version = checkpoint["channel_versions"][k] + checkpoint["versions_seen"][INTERRUPT][k] = version + self.checkpointer.put(config, create_checkpoint(checkpoint, channels)) + + async def aupdate_state( + self, config: RunnableConfig, values: dict[str, Any] | Any + ) -> None: + if not self.checkpointer: + raise ValueError("No checkpointer set") + + values = ( + {self.snapshot_channels: values} + if isinstance(self.snapshot_channels, str) + else values + ) checkpoint = await self.checkpointer.aget(config) - checkpoint = checkpoint or empty_checkpoint() + checkpoint = copy_checkpoint(checkpoint) if checkpoint else empty_checkpoint() async with AsyncChannelsManager(self.channels, checkpoint) as channels: for k, v in values.items(): channels[k].update([v]) checkpoint["channel_versions"][k] += 1 + for k in self.snapshot_channels or self.channels: + version = checkpoint["channel_versions"][k] + checkpoint["versions_seen"][INTERRUPT][k] = version await self.checkpointer.aput( config, create_checkpoint(checkpoint, channels) ) + def _defaults( + self, + debug: Optional[bool] = None, + input_keys: Optional[Union[str, Sequence[str]]] = None, + output_keys: Optional[Union[str, Sequence[str]]] = None, + interrupt_before_nodes: Optional[Sequence[str]] = None, + interrupt_after_nodes: Optional[Sequence[str]] = None, + ) -> tuple[ + bool, + Union[str, Sequence[str]], + Union[str, Sequence[str]], + Optional[Sequence[str]], + Optional[Sequence[str]], + ]: + debug = debug if debug is not None else self.debug + if output_keys is None: + output_keys = [chan for chan in self.channels if chan not in self.hidden] + else: + validate_keys(output_keys, self.channels) + if input_keys is None: + input_keys = self.input + else: + validate_keys(input_keys, self.channels) + interrupt_before_nodes = interrupt_before_nodes or self.interrupt_before_nodes + interrupt_after_nodes = interrupt_after_nodes or self.interrupt_after_nodes + return ( + debug, + input_keys, + output_keys, + interrupt_before_nodes, + interrupt_after_nodes, + ) + def _transform( self, input: Iterator[Union[dict[str, Any], Any]], run_manager: CallbackManagerForChainRun, config: RunnableConfig, *, + debug: Optional[bool] = None, input_keys: Optional[Union[str, Sequence[str]]] = None, output_keys: Optional[Union[str, Sequence[str]]] = None, - interrupt: Optional[Sequence[str]] = None, + interrupt_before_nodes: Optional[Sequence[str]] = None, + interrupt_after_nodes: Optional[Sequence[str]] = None, ) -> Iterator[Union[dict[str, Any], Any]]: try: if config["recursion_limit"] < 1: raise ValueError("recursion_limit must be at least 1") # assign defaults - if output_keys is None: - output_keys = [ - chan for chan in self.channels if chan not in self.hidden - ] - else: - validate_keys(output_keys, self.channels) - if input_keys is None: - input_keys = self.input - else: - validate_keys(input_keys, self.channels) - interrupt = interrupt or self.interrupt + ( + debug, + input_keys, + output_keys, + interrupt_before_nodes, + interrupt_after_nodes, + ) = self._defaults( + debug, + input_keys, + output_keys, + interrupt_before_nodes, + interrupt_after_nodes, + ) # copy nodes to ignore mutations during execution processes = {**self.nodes} # get checkpoint from saver, or create an empty one @@ -362,7 +445,7 @@ class Pregel( w for c in input for w in map_input(input_keys, c) ): # discard any unfinished tasks from previous checkpoint - _prepare_next_tasks(checkpoint, processes, channels) + checkpoint, _ = _prepare_next_tasks(checkpoint, processes, channels) # apply input writes _apply_writes( checkpoint, @@ -380,7 +463,9 @@ class Pregel( # channels are guaranteed to be immutable for the duration of the step, # with channel updates applied only at the transition between steps for step in range(config["recursion_limit"] + 1): - next_tasks = _prepare_next_tasks(checkpoint, processes, channels) + checkpoint, next_tasks = _prepare_next_tasks( + checkpoint, processes, channels + ) # if no more tasks, we're done if not next_tasks: @@ -392,7 +477,7 @@ class Pregel( "by setting the `recursion_limit` config key." ) - if self.debug: + if debug: print_step_start(step, next_tasks) # collect all writes to channels, without applying them yet @@ -430,15 +515,15 @@ class Pregel( timeout=self.step_timeout, ) - # interrupt on failure or timeout - _interrupt_or_proceed(done, inflight, step) + # panic on failure or timeout + _panic_or_proceed(done, inflight, step) # apply writes to channels _apply_writes( checkpoint, channels, pending_writes, config, step + 1 ) - if self.debug: + if debug: print_checkpoint(step, channels) # yield current value and checkpoint view @@ -449,22 +534,37 @@ class Pregel( # if view was updated, apply writes to channels _apply_writes_from_view(checkpoint, channels, step_output) + # with previous step's checkpoint + if do_interrupt_before := _should_interrupt( + checkpoint, + interrupt_before_nodes, + self.snapshot_channels_list, + pending_writes, + ): + break + # save end of step checkpoint - if ( - self.checkpointer is not None - and self.checkpointer.at == CheckpointAt.END_OF_STEP + if self.checkpointer is not None and ( + self.checkpointer.at == CheckpointAt.END_OF_STEP + or interrupt_before_nodes ): checkpoint = create_checkpoint(checkpoint, channels) self.checkpointer.put(config, checkpoint) - # interrupt if any channel written to is in interrupt list - if any(chan for chan, _ in pending_writes if chan in interrupt): + # with this step's checkpoint, + if _should_interrupt( + checkpoint, + interrupt_after_nodes, + self.snapshot_channels_list, + pending_writes, + ): break # save end of run checkpoint if ( self.checkpointer is not None and self.checkpointer.at == CheckpointAt.END_OF_RUN + and not do_interrupt_before ): checkpoint = create_checkpoint(checkpoint, channels) self.checkpointer.put(config, checkpoint) @@ -482,9 +582,11 @@ class Pregel( run_manager: AsyncCallbackManagerForChainRun, config: RunnableConfig, *, + debug: Optional[bool] = None, input_keys: Optional[Union[str, Sequence[str]]] = None, output_keys: Optional[Union[str, Sequence[str]]] = None, - interrupt: Optional[Sequence[str]] = None, + interrupt_before_nodes: Optional[Sequence[str]] = None, + interrupt_after_nodes: Optional[Sequence[str]] = None, ) -> AsyncIterator[Union[dict[str, Any], Any]]: try: if config["recursion_limit"] < 1: @@ -499,17 +601,19 @@ class Pregel( None, ) # assign defaults - if output_keys is None: - output_keys = [ - chan for chan in self.channels if chan not in self.hidden - ] - else: - validate_keys(output_keys, self.channels) - if input_keys is None: - input_keys = self.input - else: - validate_keys(input_keys, self.channels) - interrupt = interrupt or self.interrupt + ( + debug, + input_keys, + output_keys, + interrupt_before_nodes, + interrupt_after_nodes, + ) = self._defaults( + debug, + input_keys, + output_keys, + interrupt_before_nodes, + interrupt_after_nodes, + ) # copy nodes to ignore mutations during execution processes = {**self.nodes} # get checkpoint from saver, or create an empty one @@ -524,7 +628,7 @@ class Pregel( [w async for c in input for w in map_input(input_keys, c)] ): # discard any unfinished tasks from previous checkpoint - _prepare_next_tasks(checkpoint, processes, channels) + checkpoint, _ = _prepare_next_tasks(checkpoint, processes, channels) # apply input writes _apply_writes( checkpoint, @@ -542,7 +646,9 @@ class Pregel( # channels are guaranteed to be immutable for the duration of the step, # channel updates being applied only at the transition between steps for step in range(config["recursion_limit"] + 1): - next_tasks = _prepare_next_tasks(checkpoint, processes, channels) + checkpoint, next_tasks = _prepare_next_tasks( + checkpoint, processes, channels + ) # if no more tasks, we're done if not next_tasks: @@ -554,7 +660,7 @@ class Pregel( "by setting the `recursion_limit` config key." ) - if self.debug: + if debug: print_step_start(step, next_tasks) # collect all writes to channels, without applying them yet @@ -599,15 +705,15 @@ class Pregel( timeout=self.step_timeout, ) - # interrupt on failure or timeout - _interrupt_or_proceed(done, inflight, step) + # panic on failure or timeout + _panic_or_proceed(done, inflight, step) # apply writes to channels _apply_writes( checkpoint, channels, pending_writes, config, step + 1 ) - if self.debug: + if debug: print_checkpoint(step, channels) # yield current value and checkpoint view @@ -618,6 +724,15 @@ class Pregel( # if view was updated, apply writes to channels _apply_writes_from_view(checkpoint, channels, step_output) + # with previous step's checkpoint + if do_interrupt_before := _should_interrupt( + checkpoint, + interrupt_before_nodes, + self.snapshot_channels_list, + pending_writes, + ): + break + # save end of step checkpoint if ( self.checkpointer is not None @@ -626,14 +741,20 @@ class Pregel( checkpoint = create_checkpoint(checkpoint, channels) await self.checkpointer.aput(config, checkpoint) - # interrupt if any channel written to is in interrupt list - if any(chan for chan, _ in pending_writes if chan in interrupt): + # with this step's checkpoint + if _should_interrupt( + checkpoint, + interrupt_after_nodes, + self.snapshot_channels_list, + pending_writes, + ): break # save end of run checkpoint if ( self.checkpointer is not None and self.checkpointer.at == CheckpointAt.END_OF_RUN + and not do_interrupt_before ): checkpoint = create_checkpoint(checkpoint, channels) await self.checkpointer.aput(config, checkpoint) @@ -762,7 +883,7 @@ class Pregel( yield chunk -def _interrupt_or_proceed( +def _panic_or_proceed( done: Union[set[concurrent.futures.Future[Any]], set[asyncio.Task[Any]]], inflight: Union[set[concurrent.futures.Future[Any]], set[asyncio.Task[Any]]], step: int, @@ -786,6 +907,24 @@ def _interrupt_or_proceed( raise TimeoutError(f"Timed out at step {step}") +def _should_interrupt( + checkpoint: Checkpoint, + interrupt_nodes: Sequence[str], + snapshot_channels: Sequence[str], + pending_writes: Sequence[tuple[str, Any]], +) -> bool: + return ( + # interrupt if any of snapshopt_channels has been updated since last interrupt + any( + checkpoint["channel_versions"][chan] + > checkpoint["versions_seen"][INTERRUPT][chan] + for chan in snapshot_channels + ) + # and any channel written to is in interrupt_nodes list + and any(chan for chan, _ in pending_writes if chan in interrupt_nodes) + ) + + def _read_channel( channels: Mapping[str, BaseChannel], chan: str, catch: bool = True ) -> Any: @@ -840,6 +979,7 @@ def _apply_writes( def _apply_writes_from_view( checkpoint: Checkpoint, channels: Mapping[str, BaseChannel], values: dict[str, Any] ) -> None: + # Apply writes to channels for chan, value in values.items(): if value == _read_channel(channels, chan): continue @@ -857,7 +997,8 @@ def _prepare_next_tasks( processes: Mapping[str, Union[ChannelInvoke, ChannelBatch]], channels: Mapping[str, BaseChannel], update_seen: bool = True, -) -> list[tuple[Runnable, Any, str]]: +) -> tuple[Checkpoint, list[tuple[Runnable, Any, str]]]: + checkpoint = copy_checkpoint(checkpoint) if update_seen else checkpoint tasks: list[tuple[Runnable, Any, str]] = [] # Check if any processes should be run in next step # If so, prepare the values to be passed to them @@ -916,8 +1057,7 @@ def _prepare_next_tasks( tasks.append((proc, val, name)) if update_seen: seen[proc.channel] = checkpoint["channel_versions"][proc.channel] - - return tasks + return checkpoint, tasks async def _aconsume(iterator: AsyncIterator[Any]) -> None: diff --git a/langgraph/pregel/validate.py b/langgraph/pregel/validate.py index 8283d093c..498d96aa2 100644 --- a/langgraph/pregel/validate.py +++ b/langgraph/pregel/validate.py @@ -2,6 +2,7 @@ from typing import Any, Mapping, Sequence, Union from langgraph.channels.base import BaseChannel from langgraph.channels.last_value import LastValue +from langgraph.constants import INTERRUPT from langgraph.pregel.read import ChannelBatch, ChannelInvoke from langgraph.pregel.reserved import ReservedChannels @@ -12,10 +13,13 @@ def validate_graph( input: Union[str, Sequence[str]], output: Union[str, Sequence[str]], hidden: Sequence[str], - interrupt: Sequence[str], + interrupt_after: Sequence[str], + interrupt_before: Sequence[str], ) -> None: subscribed_channels = set[str]() - for node in nodes.values(): + for name, node in nodes.items(): + if name == INTERRUPT: + raise ValueError(f"Node name {INTERRUPT} is reserved") if isinstance(node, ChannelInvoke): subscribed_channels.update(node.channels.values()) elif isinstance(node, ChannelBatch): @@ -56,7 +60,8 @@ def validate_graph( channels[chan] = LastValue(Any) # type: ignore[arg-type] validate_keys(hidden, channels) - validate_keys(interrupt, channels) + validate_keys(interrupt_after, channels) + validate_keys(interrupt_before, channels) def validate_keys( diff --git a/poetry.lock b/poetry.lock index 0c5dc41ef..46889d252 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.7.1 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.6.1 and should not be changed by hand. [[package]] name = "aiohttp" @@ -2520,13 +2520,13 @@ dev = ["pre-commit", "pytest-asyncio", "tox"] [[package]] name = "pytest-watcher" -version = "0.3.5" +version = "0.4.1" description = "Automatically rerun your tests on file modifications" optional = false python-versions = ">=3.7.0,<4.0.0" files = [ - {file = "pytest_watcher-0.3.5-py3-none-any.whl", hash = "sha256:af00ca52c7be22dc34c0fd3d7ffef99057207a73b05dc5161fe3b2fe91f58130"}, - {file = "pytest_watcher-0.3.5.tar.gz", hash = "sha256:8896152460ba2b1a8200c12117c6611008ec96c8b2d811f0a05ab8a82b043ff8"}, + {file = "pytest_watcher-0.4.1-py3-none-any.whl", hash = "sha256:29435669cb0124fb32d6de649fe9b1350f6dac94176313fff559ee4c2a66fd6e"}, + {file = "pytest_watcher-0.4.1.tar.gz", hash = "sha256:5a793c4c883e3a55ab2abbfa3a8cd6fa6495b3767d5f6644052cc5f3236f511a"}, ] [package.dependencies] @@ -2635,7 +2635,6 @@ files = [ {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"}, {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"}, {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"}, - {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"}, {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"}, {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"}, {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"}, @@ -3760,4 +3759,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p [metadata] lock-version = "2.0" python-versions = ">=3.9.0,<4.0" -content-hash = "0e7777d77d3b34acbfdead224a2b5c5e65ecbf890c57c29bf43f5ebff09d4c0d" +content-hash = "2d35e923bf3902e0e11a305f58d17b0efc3fbb444dff8d6cb92e070a993115c9" diff --git a/pyproject.toml b/pyproject.toml index 9e2d94a0b..53d60b19d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,7 @@ pytest-asyncio = "^0.20.3" pytest-mock = "^3.10.0" syrupy = "^4.0.2" httpx = "^0.26.0" -pytest-watcher = "^0.3.4" +pytest-watcher = "^0.4.1" langchain = "^0.1.0" aiosqlite = "^0.19.0" grandalf = "^0.8" diff --git a/tests/memory_assert.py b/tests/memory_assert.py new file mode 100644 index 000000000..4362cf939 --- /dev/null +++ b/tests/memory_assert.py @@ -0,0 +1,19 @@ +from langchain_core.pydantic_v1 import Field + +from langgraph.checkpoint.base import Checkpoint, CheckpointAt, copy_checkpoint +from langgraph.checkpoint.memory import MemorySaver + + +class MemorySaverAssertImmutable(MemorySaver): + storage_for_copies: dict[str, Checkpoint] = Field(default_factory=dict) + + at = CheckpointAt.END_OF_STEP + + def put(self, config: dict, checkpoint: dict) -> None: + # assert checkpoint hasn't been modified since last written + thread_id = config["configurable"]["thread_id"] + if saved := super().get(config): + assert self.storage_for_copies[thread_id] == saved + self.storage_for_copies[thread_id] = copy_checkpoint(checkpoint) + # call super to write checkpoint + super().put(config, checkpoint) diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 6ae5f5ee9..b943c4506 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -16,7 +16,6 @@ from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.channels.context import Context from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic -from langgraph.checkpoint.memory import MemorySaver from langgraph.checkpoint.sqlite import SqliteSaver from langgraph.graph import END, Graph from langgraph.graph.message import MessageGraph @@ -28,6 +27,7 @@ from langgraph.prebuilt.chat_agent_executor import ( from langgraph.prebuilt.tool_executor import ToolExecutor from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.reserved import ReservedChannels +from tests.memory_assert import MemorySaverAssertImmutable def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: @@ -254,9 +254,11 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") - memory = MemorySaver() + memory = MemorySaverAssertImmutable() app = Pregel( - nodes={"one": one, "two": two}, checkpointer=memory, interrupt=["inbox"] + nodes={"one": one, "two": two}, + checkpointer=memory, + interrupt_after_nodes=["inbox"], ) # start execution, stop at inbox @@ -447,7 +449,7 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None: | raise_if_above_10 ) - memory = MemorySaver() + memory = MemorySaverAssertImmutable() app = Pregel( nodes={"one": one}, @@ -899,10 +901,10 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: }, ] - # test state get/update methods + # test state get/update methods with interrupt_after app_w_interrupt = workflow.compile( - checkpointer=MemorySaver(), interrupt_after=["agent"] + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["agent"] ) config = {"configurable": {"thread_id": "1"}} @@ -1044,6 +1046,152 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: } ] + # test state get/update methods with interrupt_before + + app_w_interrupt = workflow.compile( + checkpointer=MemorySaverAssertImmutable(), interrupt_before=["tools"] + ) + config = {"configurable": {"thread_id": "2"}} + llm.i = 0 # reset the llm + + assert [ + c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config) + ] == [ + { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + } + } + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + }, + "tools": None, + }, + next=("agent:edges",), + ) + + app_w_interrupt.update_state( + config, + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + }, + }, + ) + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + }, + "tools": None, + }, + next=("agent:edges",), + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "tools": { + "input": "what is weather in sf", + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ) + ], + } + }, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ) + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + ] + + app_w_interrupt.update_state( + config, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ) + ], + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + } + }, + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "__end__": { + "input": "what is weather in sf", + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ) + ], + } + } + ] + def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None: from langchain.llms.fake import FakeStreamingListLLM @@ -1235,8 +1383,10 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None: }, ] + # test state get/update methods with interrupt_after + app_w_interrupt = workflow.compile( - checkpointer=MemorySaver(), interrupt_after=["agent"] + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["agent"] ) config = {"configurable": {"thread_id": "1"}} @@ -1345,6 +1495,121 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None: } ] + # test state get/update methods with interrupt_before + + app_w_interrupt = workflow.compile( + checkpointer=MemorySaverAssertImmutable(), + interrupt_before=["tools"], + debug=True, + ) + config = {"configurable": {"thread_id": "2"}} + llm.i = 0 # reset the llm + + assert [ + c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config) + ] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + } + } + ] + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + "intermediate_steps": [], + }, + next=("agent:edges",), + ) + + app_w_interrupt.update_state( + config, + { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ) + }, + ) + + assert app_w_interrupt.get_state(config) == StateSnapshot( + values={ + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "intermediate_steps": [], + }, + next=("agent:edges",), + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "tools": { + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ) + ], + } + }, + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + ] + + app_w_interrupt.update_state( + config, + { + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ) + }, + ) + + assert [c for c in app_w_interrupt.stream(None, config)] == [ + { + "__end__": { + "input": "what is weather in sf", + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ) + ], + } + } + ] + def test_conditional_entrypoint_graph(snapshot: SnapshotAssertion) -> None: def left(data: str) -> str: @@ -2022,7 +2287,7 @@ def test_message_graph(snapshot: SnapshotAssertion) -> None: ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaver(), interrupt_after=["agent"] + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["agent"] ) config = {"configurable": {"thread_id": "1"}} diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 6db086683..fe716c953 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -23,7 +23,6 @@ from langgraph.channels.context import Context from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.checkpoint.aiosqlite import AsyncSqliteSaver -from langgraph.checkpoint.memory import MemorySaver from langgraph.graph import END, Graph, StateGraph from langgraph.graph.message import MessageGraph from langgraph.prebuilt.chat_agent_executor import ( @@ -33,6 +32,7 @@ from langgraph.prebuilt.chat_agent_executor import ( from langgraph.prebuilt.tool_executor import ToolExecutor from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.reserved import ReservedChannels +from tests.memory_assert import MemorySaverAssertImmutable async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: @@ -261,9 +261,11 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") - memory = MemorySaver() + memory = MemorySaverAssertImmutable() app = Pregel( - nodes={"one": one, "two": two}, checkpointer=memory, interrupt=["inbox"] + nodes={"one": one, "two": two}, + checkpointer=memory, + interrupt_after_nodes=["inbox"], ) # start execution, stop at inbox @@ -461,7 +463,7 @@ async def test_invoke_checkpoint(mocker: MockerFixture) -> None: | raise_if_above_10 ) - memory = MemorySaver() + memory = MemorySaverAssertImmutable() app = Pregel( nodes={"one": one}, @@ -944,10 +946,10 @@ async def test_conditional_graph() -> None: # Check that agent (one of the nodes) has its output streamed to the logs assert "/logs/agent/streamed_output/-" in patch_paths - # test state get/update methods + # test state get/update methods with interrupt_after app_w_interrupt = workflow.compile( - checkpointer=MemorySaver(), interrupt_after=["agent"] + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["agent"] ) config = {"configurable": {"thread_id": "1"}} @@ -1092,6 +1094,155 @@ async def test_conditional_graph() -> None: } ] + # test state get/update methods with interrupt_before + + app_w_interrupt = workflow.compile( + checkpointer=MemorySaverAssertImmutable(), interrupt_before=["tools"] + ) + config = {"configurable": {"thread_id": "2"}} + llm.i = 0 + + assert [ + c + async for c in app_w_interrupt.astream( + {"input": "what is weather in sf"}, config + ) + ] == [ + { + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + } + } + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent": { + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + }, + "tools": None, + }, + next=("agent:edges",), + ) + + await app_w_interrupt.aupdate_state( + config, + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + }, + }, + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "input": "what is weather in sf", + }, + "tools": None, + }, + next=("agent:edges",), + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": { + "input": "what is weather in sf", + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ) + ], + } + }, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ) + ], + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + ] + + await app_w_interrupt.aupdate_state( + config, + { + "agent": { + "input": "what is weather in sf", + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ) + ], + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + } + }, + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "__end__": { + "input": "what is weather in sf", + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ) + ], + } + } + ] + async def test_conditional_graph_state() -> None: from langchain.llms.fake import FakeStreamingListLLM @@ -1278,8 +1429,10 @@ async def test_conditional_graph_state() -> None: }, ] + # test state get/update methods with interrupt_after + app_w_interrupt = workflow.compile( - checkpointer=MemorySaver(), interrupt_after=["agent"] + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["agent"] ) config = {"configurable": {"thread_id": "1"}} @@ -1391,6 +1544,122 @@ async def test_conditional_graph_state() -> None: } ] + # test state get/update methods with interrupt_before + + app_w_interrupt = workflow.compile( + checkpointer=MemorySaverAssertImmutable(), interrupt_before=["tools"] + ) + config = {"configurable": {"thread_id": "2"}} + llm.i = 0 # reset the llm + + assert [ + c + async for c in app_w_interrupt.astream( + {"input": "what is weather in sf"}, config + ) + ] == [ + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + } + } + ] + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", tool_input="query", log="tool:search_api:query" + ), + "intermediate_steps": [], + }, + next=("agent:edges",), + ) + + await app_w_interrupt.aupdate_state( + config, + { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ) + }, + ) + + assert await app_w_interrupt.aget_state(config) == StateSnapshot( + values={ + "input": "what is weather in sf", + "agent_outcome": AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "intermediate_steps": [], + }, + next=("agent:edges",), + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "tools": { + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ) + ], + } + }, + { + "agent": { + "agent_outcome": AgentAction( + tool="search_api", + tool_input="another", + log="tool:search_api:another", + ), + } + }, + ] + + await app_w_interrupt.aupdate_state( + config, + { + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ) + }, + ) + + assert [c async for c in app_w_interrupt.astream(None, config)] == [ + { + "__end__": { + "input": "what is weather in sf", + "agent_outcome": AgentFinish( + return_values={"answer": "a really nice answer"}, + log="finish:a really nice answer", + ), + "intermediate_steps": [ + ( + AgentAction( + tool="search_api", + tool_input="query", + log="tool:search_api:a different query", + ), + "result for query", + ) + ], + } + } + ] + async def test_conditional_entrypoint_graph() -> None: async def left(data: str) -> str: @@ -2051,7 +2320,7 @@ async def test_message_graph() -> None: ] app_w_interrupt = workflow.compile( - checkpointer=MemorySaver(), interrupt_after=["agent"] + checkpointer=MemorySaverAssertImmutable(), interrupt_after=["agent"] ) config = {"configurable": {"thread_id": "1"}} From 27cf03a444bd818a837626d9730bf0bc3bdde6fe Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sat, 24 Feb 2024 19:46:29 -0800 Subject: [PATCH 5/8] Lint --- langgraph/graph/graph.py | 2 +- langgraph/graph/state.py | 5 ++--- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 53331d2af..b4ce617a8 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -14,7 +14,7 @@ from langchain_core.runnables.graph import Graph as RunnableGraph from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.checkpoint import BaseCheckpointSaver -from langgraph.pregel import Channel, Pregel, StateSnapshot +from langgraph.pregel import Channel, Pregel logger = logging.getLogger(__name__) diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index b1f048074..5c6b1d556 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -1,11 +1,10 @@ from collections import defaultdict from functools import partial from inspect import signature -from typing import Any, Optional, Sequence, Type, Union +from typing import Any, Optional, Sequence, Type from langchain_core.runnables import RunnableLambda from langchain_core.runnables.base import RunnableLike -from langchain_core.runnables.config import RunnableConfig from langgraph.channels.any_value import AnyValue from langgraph.channels.base import BaseChannel, InvalidUpdateError @@ -14,7 +13,7 @@ from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.checkpoint import BaseCheckpointSaver from langgraph.graph.graph import END, START, CompiledGraph, Graph -from langgraph.pregel import Channel, StateSnapshot +from langgraph.pregel import Channel from langgraph.pregel.read import ChannelInvoke from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry From 60097c5f1d23e6682d0d9240d1f18fdc1da86091 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sat, 24 Feb 2024 19:47:06 -0800 Subject: [PATCH 6/8] Remove class --- langgraph/graph/state.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index 5c6b1d556..76362979d 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -135,7 +135,7 @@ class StateGraph(Graph): else: raise ValueError("No entry point set") - return CompiledStateGraph( + return CompiledGraph( graph=self, nodes=nodes, channels={ @@ -205,7 +205,3 @@ def _is_field_binop(typ: Type[Any]) -> Optional[BinaryOperatorAggregate]: ): return BinaryOperatorAggregate(typ, meta[0]) return None - - -class CompiledStateGraph(CompiledGraph): - graph: StateGraph From 28aee50a9fd2f13ff60cf9aa7b4e635193850b9c Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sat, 24 Feb 2024 19:52:43 -0800 Subject: [PATCH 7/8] Lint --- langgraph/pregel/__init__.py | 9 ++++----- langgraph/pregel/reserved.py | 3 +++ 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 76df9f591..aef1a69d0 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -65,7 +65,7 @@ from langgraph.pregel.debug import print_checkpoint, print_step_start from langgraph.pregel.io import map_input, map_output from langgraph.pregel.log import logger from langgraph.pregel.read import ChannelBatch, ChannelInvoke -from langgraph.pregel.reserved import ReservedChannels +from langgraph.pregel.reserved import AllReservedChannels, ReservedChannels from langgraph.pregel.validate import validate_graph, validate_keys from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry @@ -274,7 +274,8 @@ class Pregel( return ( [self.snapshot_channels] if isinstance(self.snapshot_channels, str) - else self.snapshot_channels or list(self.channels.keys()) + else self.snapshot_channels + or [k for k in self.channels if k not in AllReservedChannels] ) def get_state(self, config: RunnableConfig) -> StateSnapshot: @@ -291,7 +292,6 @@ class Pregel( k: _read_channel(channels, k) for k in channels if k in self.snapshot_channels_list - and k not in [k.value for k in ReservedChannels] } return StateSnapshot( values[self.snapshot_channels] @@ -314,7 +314,6 @@ class Pregel( k: _read_channel(channels, k) for k in channels if k in self.snapshot_channels_list - and k not in [k.value for k in ReservedChannels] } return StateSnapshot( values[self.snapshot_channels] @@ -947,7 +946,7 @@ def _apply_writes( pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list) # Group writes by channel for chan, val in pending_writes: - if chan in [c.value for c in ReservedChannels]: + if chan in AllReservedChannels: raise ValueError(f"Can't write to reserved channel {chan}") pending_writes_by_channel[chan].append(val) diff --git a/langgraph/pregel/reserved.py b/langgraph/pregel/reserved.py index b6e66b945..2fad3b348 100644 --- a/langgraph/pregel/reserved.py +++ b/langgraph/pregel/reserved.py @@ -6,3 +6,6 @@ class ReservedChannels(StrEnum): is_last_step = "is_last_step" """A channel that is True if the current step is the last step, False otherwise.""" + + +AllReservedChannels = {channel.value for channel in ReservedChannels} From 515a512f7e31651accf3e2e71dd9d3be91aaa0a6 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sat, 24 Feb 2024 19:56:00 -0800 Subject: [PATCH 8/8] Expose additional kwargs --- langgraph/pregel/__init__.py | 69 ++++++++++++++++++++++-------------- 1 file changed, 42 insertions(+), 27 deletions(-) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index aef1a69d0..5f3064382 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -370,11 +370,12 @@ class Pregel( def _defaults( self, - debug: Optional[bool] = None, + *, input_keys: Optional[Union[str, Sequence[str]]] = None, output_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before_nodes: Optional[Sequence[str]] = None, interrupt_after_nodes: Optional[Sequence[str]] = None, + debug: Optional[bool] = None, ) -> tuple[ bool, Union[str, Sequence[str]], @@ -406,12 +407,7 @@ class Pregel( input: Iterator[Union[dict[str, Any], Any]], run_manager: CallbackManagerForChainRun, config: RunnableConfig, - *, - debug: Optional[bool] = None, - input_keys: Optional[Union[str, Sequence[str]]] = None, - output_keys: Optional[Union[str, Sequence[str]]] = None, - interrupt_before_nodes: Optional[Sequence[str]] = None, - interrupt_after_nodes: Optional[Sequence[str]] = None, + **kwargs: Any, ) -> Iterator[Union[dict[str, Any], Any]]: try: if config["recursion_limit"] < 1: @@ -423,13 +419,7 @@ class Pregel( output_keys, interrupt_before_nodes, interrupt_after_nodes, - ) = self._defaults( - debug, - input_keys, - output_keys, - interrupt_before_nodes, - interrupt_after_nodes, - ) + ) = self._defaults(**kwargs) # copy nodes to ignore mutations during execution processes = {**self.nodes} # get checkpoint from saver, or create an empty one @@ -580,12 +570,7 @@ class Pregel( input: AsyncIterator[Union[dict[str, Any], Any]], run_manager: AsyncCallbackManagerForChainRun, config: RunnableConfig, - *, - debug: Optional[bool] = None, - input_keys: Optional[Union[str, Sequence[str]]] = None, - output_keys: Optional[Union[str, Sequence[str]]] = None, - interrupt_before_nodes: Optional[Sequence[str]] = None, - interrupt_after_nodes: Optional[Sequence[str]] = None, + **kwargs: Any, ) -> AsyncIterator[Union[dict[str, Any], Any]]: try: if config["recursion_limit"] < 1: @@ -606,13 +591,7 @@ class Pregel( output_keys, interrupt_before_nodes, interrupt_after_nodes, - ) = self._defaults( - debug, - input_keys, - output_keys, - interrupt_before_nodes, - interrupt_after_nodes, - ) + ) = self._defaults(**kwargs) # copy nodes to ignore mutations during execution processes = {**self.nodes} # get checkpoint from saver, or create an empty one @@ -772,6 +751,9 @@ class Pregel( *, output_keys: Optional[Union[str, Sequence[str]]] = None, input_keys: Optional[Union[str, Sequence[str]]] = None, + interrupt_before_nodes: Optional[Sequence[str]] = None, + interrupt_after_nodes: Optional[Sequence[str]] = None, + debug: Optional[bool] = None, **kwargs: Any, ) -> Union[dict[str, Any], Any]: latest: Union[dict[str, Any], Any] = None @@ -780,6 +762,9 @@ class Pregel( config, output_keys=output_keys if output_keys is not None else self.output, input_keys=input_keys, + interrupt_before_nodes=interrupt_before_nodes, + interrupt_after_nodes=interrupt_after_nodes, + debug=debug, **kwargs, ): latest = chunk @@ -792,6 +777,9 @@ class Pregel( *, output_keys: Optional[Union[str, Sequence[str]]] = None, input_keys: Optional[Union[str, Sequence[str]]] = None, + interrupt_before_nodes: Optional[Sequence[str]] = None, + interrupt_after_nodes: Optional[Sequence[str]] = None, + debug: Optional[bool] = None, **kwargs: Any, ) -> Iterator[Union[dict[str, Any], Any]]: return self.transform( @@ -799,6 +787,9 @@ class Pregel( config, output_keys=output_keys, input_keys=input_keys, + interrupt_before_nodes=interrupt_before_nodes, + interrupt_after_nodes=interrupt_after_nodes, + debug=debug, **kwargs, ) @@ -809,6 +800,9 @@ class Pregel( *, output_keys: Optional[Union[str, Sequence[str]]] = None, input_keys: Optional[Union[str, Sequence[str]]] = None, + interrupt_before_nodes: Optional[Sequence[str]] = None, + interrupt_after_nodes: Optional[Sequence[str]] = None, + debug: Optional[bool] = None, **kwargs: Any, ) -> Iterator[Union[dict[str, Any], Any]]: for chunk in self._transform_stream_with_config( @@ -817,6 +811,9 @@ class Pregel( config, output_keys=output_keys, input_keys=input_keys, + interrupt_before_nodes=interrupt_before_nodes, + interrupt_after_nodes=interrupt_after_nodes, + debug=debug, **kwargs, ): yield chunk @@ -828,6 +825,9 @@ class Pregel( *, output_keys: Optional[Union[str, Sequence[str]]] = None, input_keys: Optional[Union[str, Sequence[str]]] = None, + interrupt_before_nodes: Optional[Sequence[str]] = None, + interrupt_after_nodes: Optional[Sequence[str]] = None, + debug: Optional[bool] = None, **kwargs: Any, ) -> Union[dict[str, Any], Any]: latest: Union[dict[str, Any], Any] = None @@ -836,6 +836,9 @@ class Pregel( config, output_keys=output_keys if output_keys is not None else self.output, input_keys=input_keys, + interrupt_before_nodes=interrupt_before_nodes, + interrupt_after_nodes=interrupt_after_nodes, + debug=debug, **kwargs, ): latest = chunk @@ -848,6 +851,9 @@ class Pregel( *, output_keys: Optional[Union[str, Sequence[str]]] = None, input_keys: Optional[Union[str, Sequence[str]]] = None, + interrupt_before_nodes: Optional[Sequence[str]] = None, + interrupt_after_nodes: Optional[Sequence[str]] = None, + debug: Optional[bool] = None, **kwargs: Any, ) -> AsyncIterator[Union[dict[str, Any], Any]]: async def input_stream() -> AsyncIterator[Union[dict[str, Any], Any]]: @@ -858,6 +864,9 @@ class Pregel( config, output_keys=output_keys, input_keys=input_keys, + interrupt_before_nodes=interrupt_before_nodes, + interrupt_after_nodes=interrupt_after_nodes, + debug=debug, **kwargs, ): yield chunk @@ -869,6 +878,9 @@ class Pregel( *, output_keys: Optional[Union[str, Sequence[str]]] = None, input_keys: Optional[Union[str, Sequence[str]]] = None, + interrupt_before_nodes: Optional[Sequence[str]] = None, + interrupt_after_nodes: Optional[Sequence[str]] = None, + debug: Optional[bool] = None, **kwargs: Any, ) -> AsyncIterator[Union[dict[str, Any], Any]]: async for chunk in self._atransform_stream_with_config( @@ -877,6 +889,9 @@ class Pregel( config, output_keys=output_keys, input_keys=input_keys, + interrupt_before_nodes=interrupt_before_nodes, + interrupt_after_nodes=interrupt_after_nodes, + debug=debug, **kwargs, ): yield chunk