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]: