diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index a855597c9..b21b6f0b3 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -247,10 +247,9 @@ class Graph: 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], - snapshot_channels=list(self.nodes), + input_channels=f"{self.entry_point}:inbox" if self.entry_point else START, + output_channels=END, + stream_channels=list(self.nodes), checkpointer=checkpointer, interrupt_before_nodes=[f"{node}:inbox" for node in interrupt_before], interrupt_after_nodes=interrupt_after, diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index f7b98db53..b61581736 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -186,13 +186,10 @@ class StateGraph(Graph): **waiting_edge_channels, END: LastValue(self.schema), }, - input=f"{START}:inbox", - output=END, - hidden=[f"{node}:inbox" for node in self.nodes] - + [START] - + state_keys - + [key for key, _, _ in waiting_edges], - snapshot_channels=state_keys_read, + input_channels=f"{START}:inbox", + stream_mode="updates", + output_channels=END, + stream_channels=state_keys_read, checkpointer=checkpointer, interrupt_before_nodes=[f"{node}:inbox" for node in interrupt_before], interrupt_after_nodes=interrupt_after, diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index e66173293..8dc4bedcf 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -61,7 +61,7 @@ from langgraph.checkpoint.base import ( ) 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.io import map_input, map_output_updates, map_output_values from langgraph.pregel.log import logger from langgraph.pregel.read import ChannelInvoke from langgraph.pregel.reserved import AllReservedChannels, ReservedChannels @@ -156,6 +156,9 @@ class Channel: ) +StreamMode = Literal["values", "updates"] + + class StateSnapshot(NamedTuple): values: dict[str, Any] | Any """Current values of channels""" @@ -174,20 +177,19 @@ class Pregel( channels: Mapping[str, BaseChannel] = Field(default_factory=dict) - # TODO Rename to `output_channels` - output: Union[str, Sequence[str]] = "output" + stream_mode: StreamMode = "values" - # TODO Replace with `stream_channels` - hidden: Sequence[str] = Field(default_factory=list) + output_channels: Union[str, Sequence[str]] = "output" + """Channels to output, defaults to channel named 'output'.""" - snapshot_channels: Union[str, Sequence[str]] = Field(default_factory=list) + stream_channels: Optional[Union[str, Sequence[str]]] = None + """Channels to stream, defaults to all channels not in reserved channels""" 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" + input_channels: Union[str, Sequence[str]] = "input" step_timeout: Optional[float] = None @@ -205,9 +207,9 @@ class Pregel( validate_graph( values["nodes"], values["channels"], - values["input"], - values["output"], - values["hidden"], + values["input_channels"], + values["output_channels"], + values["stream_channels"], values["interrupt_after_nodes"], values["interrupt_before_nodes"], ) @@ -234,45 +236,45 @@ class Pregel( @property def InputType(self) -> Any: - if isinstance(self.input, str): - return self.channels[self.input].UpdateType + if isinstance(self.input_channels, str): + return self.channels[self.input_channels].UpdateType def get_input_schema( self, config: Optional[RunnableConfig] = None ) -> Type[BaseModel]: - if isinstance(self.input, str): + if isinstance(self.input_channels, str): return super().get_input_schema(config) else: return create_model( # type: ignore[call-overload] self.get_name("Input"), **{ k: (self.channels[k].UpdateType, None) - for k in self.input or self.channels.keys() + for k in self.input_channels or self.channels.keys() }, ) @property def OutputType(self) -> Any: - if isinstance(self.output, str): - return self.channels[self.output].ValueType + if isinstance(self.output_channels, str): + return self.channels[self.output_channels].ValueType def get_output_schema( self, config: Optional[RunnableConfig] = None ) -> Type[BaseModel]: - if isinstance(self.output, str): + if isinstance(self.output_channels, str): return super().get_output_schema(config) else: return create_model( # type: ignore[call-overload] self.get_name("Output"), - **{k: (self.channels[k].ValueType, None) for k in self.output}, + **{k: (self.channels[k].ValueType, None) for k in self.output_channels}, ) @property def snapshot_channels_list(self) -> Sequence[str]: return ( - [self.snapshot_channels] - if isinstance(self.snapshot_channels, str) - else self.snapshot_channels + [self.stream_channels] + if isinstance(self.stream_channels, str) + else self.stream_channels or [k for k in self.channels if k not in AllReservedChannels] ) @@ -293,8 +295,8 @@ class Pregel( if k in self.snapshot_channels_list } return StateSnapshot( - values[self.snapshot_channels] - if isinstance(self.snapshot_channels, str) + values[self.stream_channels] + if isinstance(self.stream_channels, str) else values, tuple(name for _, _, name in next_tasks), config, @@ -317,8 +319,8 @@ class Pregel( if k in self.snapshot_channels_list } return StateSnapshot( - values[self.snapshot_channels] - if isinstance(self.snapshot_channels, str) + values[self.stream_channels] + if isinstance(self.stream_channels, str) else values, tuple(name for _, _, name in next_tasks), config, @@ -339,8 +341,8 @@ class Pregel( if k in self.snapshot_channels_list } yield StateSnapshot( - values[self.snapshot_channels] - if isinstance(self.snapshot_channels, str) + values[self.stream_channels] + if isinstance(self.stream_channels, str) else values, tuple(name for _, _, name in next_tasks), config, @@ -364,8 +366,8 @@ class Pregel( if k in self.snapshot_channels_list } yield StateSnapshot( - values[self.snapshot_channels] - if isinstance(self.snapshot_channels, str) + values[self.stream_channels] + if isinstance(self.stream_channels, str) else values, tuple(name for _, _, name in next_tasks), config, @@ -379,8 +381,8 @@ class Pregel( raise ValueError("No checkpointer set") values = ( - {self.snapshot_channels: values} - if isinstance(self.snapshot_channels, str) + {self.stream_channels: values} + if isinstance(self.stream_channels, str) else values ) checkpoint = self.checkpointer.get(config) @@ -403,8 +405,8 @@ class Pregel( raise ValueError("No checkpointer set") values = ( - {self.snapshot_channels: values} - if isinstance(self.snapshot_channels, str) + {self.stream_channels: values} + if isinstance(self.stream_channels, str) else values ) checkpoint = await self.checkpointer.aget(config) @@ -413,7 +415,7 @@ class Pregel( for k, v in values.items(): channels[k].update([v]) checkpoint["channel_versions"][k] += 1 - for k in self.snapshot_channels or self.channels: + for k in self.stream_channels or self.channels: version = checkpoint["channel_versions"][k] checkpoint["versions_seen"][INTERRUPT][k] = version return await self.checkpointer.aput( @@ -423,6 +425,7 @@ class Pregel( def _defaults( self, *, + stream_mode: Optional[StreamMode] = None, input_keys: Optional[Union[str, Sequence[str]]] = None, output_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before_nodes: Optional[Sequence[str]] = None, @@ -430,6 +433,7 @@ class Pregel( debug: Optional[bool] = None, ) -> tuple[ bool, + StreamMode, Union[str, Sequence[str]], Union[str, Sequence[str]], Optional[Sequence[str]], @@ -437,17 +441,22 @@ class Pregel( ]: 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] + output_keys = ( + [chan for chan in self.channels] + if self.stream_channels is None + else self.stream_channels + ) else: validate_keys(output_keys, self.channels) if input_keys is None: - input_keys = self.input + input_keys = self.input_channels 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, + stream_mode if stream_mode is not None else self.stream_mode, input_keys, output_keys, interrupt_before_nodes, @@ -467,6 +476,7 @@ class Pregel( # assign defaults ( debug, + stream_mode, input_keys, output_keys, interrupt_before_nodes, @@ -584,9 +594,15 @@ class Pregel( if debug: print_checkpoint(step, channels) - # yield current value and checkpoint view - if step_output := map_output(output_keys, pending_writes, channels): - yield step_output + # yield current value or updates + if stream_mode == "values": + if step_output := map_output_values( + output_keys, pending_writes, channels + ): + yield step_output + else: + if step_output := map_output_updates(output_keys, next_tasks): + yield step_output # with previous step's checkpoint if _should_interrupt( @@ -654,6 +670,7 @@ class Pregel( # assign defaults ( debug, + stream_mode, input_keys, output_keys, interrupt_before_nodes, @@ -775,9 +792,15 @@ class Pregel( if debug: print_checkpoint(step, channels) - # yield current value and checkpoint view - if step_output := map_output(output_keys, pending_writes, channels): - yield step_output + # yield current value or updates + if stream_mode == "values": + if step_output := map_output_values( + output_keys, pending_writes, channels + ): + yield step_output + else: + if step_output := map_output_updates(output_keys, next_tasks): + yield step_output # with previous step's checkpoint if _should_interrupt( @@ -835,18 +858,21 @@ class Pregel( debug: Optional[bool] = None, **kwargs: Any, ) -> Union[dict[str, Any], Any]: - latest: Union[dict[str, Any], Any] = None + output_keys = output_keys if output_keys is not None else self.output_channels + output_is_dict = not isinstance(output_keys, str) + latest: Union[dict[str, Any], Any] = {} if output_is_dict else None for chunk in self.stream( input, config, - output_keys=output_keys if output_keys is not None else self.output, + stream_mode="values", + output_keys=output_keys, input_keys=input_keys, interrupt_before_nodes=interrupt_before_nodes, interrupt_after_nodes=interrupt_after_nodes, debug=debug, **kwargs, ): - latest = chunk + latest = {**latest, **chunk} if output_is_dict else chunk return latest def stream( @@ -854,6 +880,7 @@ class Pregel( input: Union[dict[str, Any], Any], config: Optional[RunnableConfig] = None, *, + stream_mode: Optional[StreamMode] = None, output_keys: Optional[Union[str, Sequence[str]]] = None, input_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before_nodes: Optional[Sequence[str]] = None, @@ -864,6 +891,7 @@ class Pregel( return self.transform( iter([input]), config, + stream_mode=stream_mode, output_keys=output_keys, input_keys=input_keys, interrupt_before_nodes=interrupt_before_nodes, @@ -877,6 +905,7 @@ class Pregel( input: Iterator[Union[dict[str, Any], Any]], config: Optional[RunnableConfig] = None, *, + stream_mode: Optional[StreamMode] = None, output_keys: Optional[Union[str, Sequence[str]]] = None, input_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before_nodes: Optional[Sequence[str]] = None, @@ -888,6 +917,7 @@ class Pregel( input, self._transform, config, + stream_mode=stream_mode, output_keys=output_keys, input_keys=input_keys, interrupt_before_nodes=interrupt_before_nodes, @@ -909,18 +939,21 @@ class Pregel( debug: Optional[bool] = None, **kwargs: Any, ) -> Union[dict[str, Any], Any]: - latest: Union[dict[str, Any], Any] = None + output_keys = output_keys if output_keys is not None else self.output_channels + output_is_dict = not isinstance(output_keys, str) + latest: Union[dict[str, Any], Any] = {} if output_is_dict else None async for chunk in self.astream( input, config, - output_keys=output_keys if output_keys is not None else self.output, + stream_mode="values", + output_keys=output_keys, input_keys=input_keys, interrupt_before_nodes=interrupt_before_nodes, interrupt_after_nodes=interrupt_after_nodes, debug=debug, **kwargs, ): - latest = chunk + latest = {**latest, **chunk} if output_is_dict else chunk return latest async def astream( @@ -928,6 +961,7 @@ class Pregel( input: Union[dict[str, Any], Any], config: Optional[RunnableConfig] = None, *, + stream_mode: Optional[StreamMode] = None, output_keys: Optional[Union[str, Sequence[str]]] = None, input_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before_nodes: Optional[Sequence[str]] = None, @@ -941,6 +975,7 @@ class Pregel( async for chunk in self.atransform( input_stream(), config, + stream_mode=stream_mode, output_keys=output_keys, input_keys=input_keys, interrupt_before_nodes=interrupt_before_nodes, @@ -955,6 +990,7 @@ class Pregel( input: AsyncIterator[Union[dict[str, Any], Any]], config: Optional[RunnableConfig] = None, *, + stream_mode: Optional[StreamMode] = None, output_keys: Optional[Union[str, Sequence[str]]] = None, input_keys: Optional[Union[str, Sequence[str]]] = None, interrupt_before_nodes: Optional[Sequence[str]] = None, @@ -966,6 +1002,7 @@ class Pregel( input, self._atransform, config, + stream_mode=stream_mode, output_keys=output_keys, input_keys=input_keys, interrupt_before_nodes=interrupt_before_nodes, diff --git a/langgraph/pregel/io.py b/langgraph/pregel/io.py index 829cb9242..3046db221 100644 --- a/langgraph/pregel/io.py +++ b/langgraph/pregel/io.py @@ -1,5 +1,8 @@ +from collections import deque from typing import Any, Iterator, Mapping, Optional, Sequence, Union +from langchain_core.runnables import Runnable + from langgraph.channels.base import BaseChannel, EmptyChannelError from langgraph.pregel.log import logger @@ -35,7 +38,7 @@ def map_input( logger.warning(f"Input channel {k} not found in {input_channels}") -def map_output( +def map_output_values( output_channels: Union[str, Sequence[str]], pending_writes: Sequence[tuple[str, Any]], channels: Mapping[str, BaseChannel], @@ -48,3 +51,26 @@ def map_output( if updated := {c for c, _ in pending_writes if c in output_channels}: return {chan: _read_channel(channels, chan) for chan in updated} return None + + +def map_output_updates( + output_channels: Union[str, Sequence[str]], + next_tasks: list[tuple[Runnable, Any, str, deque[tuple[str, Any]]]], +) -> Optional[dict[str, Union[Any, dict[str, Any]]]]: + """Map pending writes (a sequence of tuples (channel, value)) to output chunk.""" + if isinstance(output_channels, str): + if updated := { + node: value + for _, _, node, writes in next_tasks + for chan, value in writes + if chan == output_channels + }: + return updated + else: + if updated := { + node: {chan: value for chan, value in writes if chan in output_channels} + for _, _, node, writes in next_tasks + if any(chan in output_channels for chan, _ in writes) + }: + return updated + return None diff --git a/langgraph/pregel/validate.py b/langgraph/pregel/validate.py index 302937689..4d3531b20 100644 --- a/langgraph/pregel/validate.py +++ b/langgraph/pregel/validate.py @@ -1,4 +1,4 @@ -from typing import Any, Mapping, Sequence, Union +from typing import Any, Mapping, Optional, Sequence, Union from langgraph.channels.base import BaseChannel from langgraph.channels.last_value import LastValue @@ -10,11 +10,11 @@ from langgraph.pregel.reserved import ReservedChannels def validate_graph( nodes: Mapping[str, ChannelInvoke], channels: dict[str, BaseChannel], - input: Union[str, Sequence[str]], - output: Union[str, Sequence[str]], - hidden: Sequence[str], - interrupt_after: Sequence[str], - interrupt_before: Sequence[str], + input_channels: Union[str, Sequence[str]], + output_channels: Union[str, Sequence[str]], + stream_channels: Optional[Union[str, Sequence[str]]], + interrupt_after_nodes: Sequence[str], + interrupt_before_nodes: Sequence[str], ) -> None: subscribed_channels = set[str]() for name, node in nodes.items(): @@ -31,25 +31,27 @@ def validate_graph( if chan not in channels: channels[chan] = LastValue(Any) # type: ignore[arg-type] - if isinstance(input, str): - if input not in channels: - channels[input] = LastValue(Any) # type: ignore[arg-type] - if input not in subscribed_channels: - raise ValueError(f"Input channel {input} is not subscribed to by any node") + if isinstance(input_channels, str): + if input_channels not in channels: + channels[input_channels] = LastValue(Any) # type: ignore[arg-type] + if input_channels not in subscribed_channels: + raise ValueError( + f"Input channel {input_channels} is not subscribed to by any node" + ) else: - for chan in input: + for chan in input_channels: if chan not in channels: channels[chan] = LastValue(Any) # type: ignore[arg-type] - if all(chan not in subscribed_channels for chan in input): + if all(chan not in subscribed_channels for chan in input_channels): raise ValueError( - f"None of the input channels {input} are subscribed to by any node" + f"None of the input channels {input_channels} are subscribed to by any node" ) - if isinstance(output, str): - if output not in channels: - channels[output] = LastValue(Any) # type: ignore[arg-type] + if isinstance(output_channels, str): + if output_channels not in channels: + channels[output_channels] = LastValue(Any) # type: ignore[arg-type] else: - for chan in output: + for chan in output_channels: if chan not in channels: channels[chan] = LastValue(Any) # type: ignore[arg-type] @@ -57,19 +59,19 @@ def validate_graph( if chan not in channels: channels[chan] = LastValue(Any) # type: ignore[arg-type] - validate_keys(hidden, channels) - validate_keys(interrupt_after, channels) - validate_keys(interrupt_before, channels) + validate_keys(stream_channels, channels) + validate_keys(interrupt_after_nodes, channels) + validate_keys(interrupt_before_nodes, channels) def validate_keys( - keys: Union[str, Sequence[str]], + keys: Optional[Union[str, Sequence[str]]], channels: Mapping[str, BaseChannel], ) -> None: if isinstance(keys, str): if keys not in channels: raise ValueError(f"Key {keys} not in channels") - else: + elif keys is not None: for chan in keys: if chan not in channels: raise ValueError(f"Key {chan} not in channels") diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 238a83670..3a26a2c67 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -42,8 +42,8 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: "input": LastValue(int), "output": LastValue(int), }, - input="input", - output="output", + input_channels="input", + output_channels="output", ) graph = Graph() graph.add_node("add_one", add_one) @@ -86,7 +86,9 @@ def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None: | Channel.write_to("output", fixed=5, output_plus_one=lambda x: x + 1) ) - app = Pregel(nodes={"one": chain}, output=["output", "fixed", "output_plus_one"]) + app = Pregel( + nodes={"one": chain}, output_channels=["output", "fixed", "output_plus_one"] + ) assert app.input_schema.schema() == {"title": "LangGraphInput"} assert app.output_schema.schema() == { @@ -126,7 +128,7 @@ def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: nodes={ "one": chain, }, - output=["output"], + output_channels=["output"], ) assert app.input_schema.schema() == {"title": "LangGraphInput"} @@ -146,8 +148,8 @@ def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None: nodes={ "one": chain, }, - input=["input"], - output=["output"], + input_channels=["input"], + output_channels=["output"], ) assert app.input_schema.schema() == { @@ -198,13 +200,9 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert values == { "add_one_more": 4, } - elif step == 3: - assert values == { - "__end__": 4, - } else: assert 0, f"{step}:{values}" - assert step == 3 + assert step == 2 def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: @@ -271,13 +269,27 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: app = Pregel( nodes={"one": one, "two": two}, channels={"inbox": Topic(int)}, - input=["input", "inbox"], + input_channels=["input", "inbox"], ) + # [12 + 1, 2 + 1 + 1] + assert [ + *app.stream( + {"input": 2, "inbox": 12}, output_keys="output", stream_mode="updates" + ) + ] == [ + {"two": 13}, + {"two": 4}, + ] assert [*app.stream({"input": 2, "inbox": 12}, output_keys="output")] == [ 13, 4, - ] # [12 + 1, 2 + 1 + 1] + ] + + assert [*app.stream({"input": 2, "inbox": 12}, stream_mode="updates")] == [ + {"one": {"inbox": 3}, "two": {"output": 13}}, + {"two": {"output": 4}}, + ] assert [*app.stream({"input": 2, "inbox": 12})] == [ {"inbox": [3], "output": 13}, {"output": 4}, @@ -616,6 +628,10 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None: app = Pregel(nodes={"one": one, "two": two}) + assert [c for c in app.stream(2, stream_mode="updates")] == [ + {"one": {"between": 3, "output": 3}}, + {"two": {"output": 4}}, + ] assert [c for c in app.stream(2)] == [{"between": 3, "output": 3}, {"output": 4}] @@ -667,7 +683,7 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: "inbox": Topic(int), "ctx": Context(an_int, typ=int), }, - output=["inbox", "output"], + output_channels=["inbox", "output"], ) assert setup.call_count == 0 @@ -881,32 +897,6 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: ), } }, - { - "__end__": { - "input": "what is weather in sf", - "intermediate_steps": [ - ( - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ), - ( - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ), - ], - "agent_outcome": AgentFinish( - return_values={"answer": "answer"}, log="finish:answer" - ), - } - }, ] # test state get/update methods with interrupt_after @@ -1040,27 +1030,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: }, ) - 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", - ) - ], - } - } - ] + assert [c for c in app_w_interrupt.stream(None, config)] == [] # test state get/update methods with interrupt_before @@ -1188,27 +1158,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: }, ) - 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", - ) - ], - } - } - ] + assert [c for c in app_w_interrupt.stream(None, config)] == [] # test re-invoke to continue with interrupt_before @@ -1333,32 +1283,6 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: ), } }, - { - "__end__": { - "input": "what is weather in sf", - "intermediate_steps": [ - ( - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ), - ( - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ), - ], - "agent_outcome": AgentFinish( - return_values={"answer": "answer"}, log="finish:answer" - ), - } - }, ] @@ -1473,6 +1397,7 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None: } assert [*app.stream({"input": "what is weather in sf"})] == [ + {"__start__": {"input": "what is weather in sf"}}, { "agent": { "agent_outcome": AgentAction( @@ -1524,32 +1449,6 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None: ), } }, - { - "__end__": { - "input": "what is weather in sf", - "intermediate_steps": [ - ( - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ), - ( - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ), - ], - "agent_outcome": AgentFinish( - return_values={"answer": "answer"}, log="finish:answer" - ), - } - }, ] # test state get/update methods with interrupt_after @@ -1562,13 +1461,14 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None: assert [ c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config) ] == [ + {"__start__": {"input": "what is weather in sf"}}, { "agent": { "agent_outcome": AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query" ), } - } + }, ] assert app_w_interrupt.get_state(config) == StateSnapshot( @@ -1644,27 +1544,7 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None: }, ) - 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", - ) - ], - } - } - ] + assert [c for c in app_w_interrupt.stream(None, config)] == [] # test state get/update methods with interrupt_before @@ -1679,13 +1559,14 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None: assert [ c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config) ] == [ + {"__start__": {"input": "what is weather in sf"}}, { "agent": { "agent_outcome": AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query" ), } - } + }, ] assert app_w_interrupt.get_state(config) == StateSnapshot( @@ -1761,27 +1642,7 @@ def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None: }, ) - 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", - ) - ], - } - } - ] + assert [c for c in app_w_interrupt.stream(None, config)] == [] def test_conditional_entrypoint_graph(snapshot: SnapshotAssertion) -> None: @@ -1822,7 +1683,6 @@ def test_conditional_entrypoint_graph(snapshot: SnapshotAssertion) -> None: assert [*app.stream("what is weather in sf")] == [ {"right": "what is weather in sf->right"}, - {"__end__": "what is weather in sf->right"}, ] @@ -1870,13 +1730,8 @@ def test_conditional_entrypoint_graph_state(snapshot: SnapshotAssertion) -> None } assert [*app.stream({"input": "what is weather in sf"})] == [ + {"__start__": {"input": "what is weather in sf"}}, {"right": {"output": "what is weather in sf->right"}}, - { - "__end__": { - "input": "what is weather in sf", - "output": "what is weather in sf->right", - } - }, ] @@ -2001,6 +1856,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: assert [ *app.stream({"messages": [HumanMessage(content="what is weather in sf")]}) ] == [ + {"__start__": {"messages": [HumanMessage(content="what is weather in sf")]}}, { "agent": { "messages": [ @@ -2071,61 +1927,6 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: } }, {"agent": {"messages": [AIMessage(content="answer")]}}, - { - "__end__": { - "messages": [ - HumanMessage(content="what is weather in sf"), - AIMessage( - content="", - additional_kwargs={ - "tool_calls": [ - { - "id": "tool_call123", - "type": "function", - "function": { - "name": "search_api", - "arguments": '"query"', - }, - } - ] - }, - ), - ToolMessage( - content="result for query", tool_call_id="tool_call123" - ), - AIMessage( - content="", - additional_kwargs={ - "tool_calls": [ - { - "id": "tool_call234", - "type": "function", - "function": { - "name": "search_api", - "arguments": '"another"', - }, - }, - { - "id": "tool_call567", - "type": "function", - "function": { - "name": "search_api", - "arguments": '"a third one"', - }, - }, - ] - }, - ), - ToolMessage( - content="result for another", tool_call_id="tool_call234" - ), - ToolMessage( - content="result for a third one", tool_call_id="tool_call567" - ), - AIMessage(content="answer"), - ] - } - }, ] @@ -2203,6 +2004,7 @@ def test_prebuilt_chat(snapshot: SnapshotAssertion) -> None: assert [ *app.stream({"messages": [HumanMessage(content="what is weather in sf")]}) ] == [ + {"__start__": {"messages": [HumanMessage(content="what is weather in sf")]}}, { "agent": { "messages": [ @@ -2248,34 +2050,6 @@ def test_prebuilt_chat(snapshot: SnapshotAssertion) -> None: } }, {"agent": {"messages": [AIMessage(content="answer")]}}, - { - "__end__": { - "messages": [ - HumanMessage(content="what is weather in sf"), - AIMessage( - content="", - additional_kwargs={ - "function_call": { - "name": "search_api", - "arguments": '"query"', - } - }, - ), - FunctionMessage(content="result for query", name="search_api"), - AIMessage( - content="", - additional_kwargs={ - "function_call": { - "name": "search_api", - "arguments": '"another"', - } - }, - ), - FunctionMessage(content="result for another", name="search_api"), - AIMessage(content="answer"), - ] - } - }, ] @@ -2456,6 +2230,14 @@ def test_message_graph( ] assert [*app.stream([HumanMessage(content="what is weather in sf")])] == [ + { + "__start__": [ + HumanMessage( + content="what is weather in sf", + id="00000000-0000-4000-8000-000000000038", + ) + ] + }, { "agent": AIMessage( content="", @@ -2489,42 +2271,6 @@ def test_message_graph( ) }, {"agent": AIMessage(content="answer", id="ai3")}, - { - "__end__": [ - HumanMessage( - content="what is weather in sf", - id="00000000-0000-4000-8000-000000000038", - ), - AIMessage( - content="", - additional_kwargs={ - "function_call": {"name": "search_api", "arguments": '"query"'} - }, - id="ai1", - ), - FunctionMessage( - content="result for query", - name="search_api", - id="00000000-0000-4000-8000-000000000051", - ), - AIMessage( - content="", - additional_kwargs={ - "function_call": { - "name": "search_api", - "arguments": '"another"', - } - }, - id="ai2", - ), - FunctionMessage( - content="result for another", - name="search_api", - id="00000000-0000-4000-8000-000000000064", - ), - AIMessage(content="answer", id="ai3"), - ] - }, ] app_w_interrupt = workflow.compile( @@ -2538,6 +2284,12 @@ def test_message_graph( HumanMessage(content="what is weather in sf"), config ) ] == [ + { + "__start__": HumanMessage( + content="what is weather in sf", + id="00000000-0000-4000-8000-000000000074", + ) + }, { "agent": AIMessage( content="", @@ -2546,7 +2298,7 @@ def test_message_graph( }, id="ai1", ) - } + }, ] assert app_w_interrupt.get_state(config) == StateSnapshot( @@ -2679,32 +2431,7 @@ def test_message_graph( config=app_w_interrupt.checkpointer.get_tuple(config).config, ) - assert [c for c in app_w_interrupt.stream(None, config)] == [ - { - "__end__": [ - HumanMessage( - content="what is weather in sf", - id="00000000-0000-4000-8000-000000000074", - ), - AIMessage( - content="", - additional_kwargs={ - "function_call": { - "name": "search_api", - "arguments": '"a different query"', - } - }, - id="ai1", - ), - FunctionMessage( - content="result for a different query", - name="search_api", - id="00000000-0000-4000-8000-000000000088", - ), - AIMessage(content="answer", id="ai2"), - ] - } - ] + assert [c for c in app_w_interrupt.stream(None, config)] == [] app_w_interrupt = workflow.compile( checkpointer=MemorySaverAssertImmutable(), interrupt_before=["action"] @@ -2718,6 +2445,12 @@ def test_message_graph( HumanMessage(content="what is weather in sf"), config ) ] == [ + { + "__start__": HumanMessage( + content="what is weather in sf", + id="00000000-0000-4000-8000-000000000099", + ) + }, { "agent": AIMessage( content="", @@ -2726,7 +2459,7 @@ def test_message_graph( }, id="ai1", ) - } + }, ] assert app_w_interrupt.get_state(config) == StateSnapshot( @@ -2859,32 +2592,7 @@ def test_message_graph( config=app_w_interrupt.checkpointer.get_tuple(config).config, ) - assert [c for c in app_w_interrupt.stream(None, config)] == [ - { - "__end__": [ - HumanMessage( - content="what is weather in sf", - id="00000000-0000-4000-8000-000000000099", - ), - AIMessage( - content="", - additional_kwargs={ - "function_call": { - "name": "search_api", - "arguments": '"a different query"', - } - }, - id="ai1", - ), - FunctionMessage( - content="result for a different query", - name="search_api", - id="00000000-0000-4000-8000-000000000116", - ), - AIMessage(content="answer", id="ai2"), - ] - } - ] + assert [c for c in app_w_interrupt.stream(None, config)] == [] def test_in_one_fan_out_out_one_graph_state() -> None: @@ -2931,19 +2639,20 @@ def test_in_one_fan_out_out_one_graph_state() -> None: } assert [*app.stream({"query": "what is weather in sf"})] == [ + {"__start__": {"query": "what is weather in sf"}}, {"rewrite_query": {"query": "query: what is weather in sf"}}, { "retriever_two": {"docs": ["doc3", "doc4"]}, "retriever_one": {"docs": ["doc1", "doc2"]}, }, {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - { - "__end__": { - "query": "query: what is weather in sf", - "answer": "doc1,doc2,doc3,doc4", - "docs": ["doc1", "doc2", "doc3", "doc4"], - } - }, + ] + + assert [*app.stream({"query": "what is weather in sf"}, stream_mode="values")] == [ + {"query": "what is weather in sf"}, + {"query": "query: what is weather in sf"}, + {"docs": ["doc1", "doc2", "doc3", "doc4"]}, + {"answer": "doc1,doc2,doc3,doc4"}, ] @@ -3037,6 +2746,7 @@ def test_in_one_fan_out_state_graph_waiting_edge() -> None: } assert [*app.stream({"query": "what is weather in sf"})] == [ + {"__start__": {"query": "what is weather in sf"}}, {"rewrite_query": {"query": "query: what is weather in sf"}}, { "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, @@ -3044,13 +2754,6 @@ def test_in_one_fan_out_state_graph_waiting_edge() -> None: }, {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - { - "__end__": { - "query": "analyzed: query: what is weather in sf", - "answer": "doc1,doc2,doc3,doc4", - "docs": ["doc1", "doc2", "doc3", "doc4"], - } - }, ] app_w_interrupt = workflow.compile( @@ -3061,6 +2764,7 @@ def test_in_one_fan_out_state_graph_waiting_edge() -> None: assert [ c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config) ] == [ + {"__start__": {"query": "what is weather in sf"}}, {"rewrite_query": {"query": "query: what is weather in sf"}}, { "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, @@ -3071,13 +2775,6 @@ def test_in_one_fan_out_state_graph_waiting_edge() -> None: assert [c for c in app_w_interrupt.stream(None, config)] == [ {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - { - "__end__": { - "query": "analyzed: query: what is weather in sf", - "answer": "doc1,doc2,doc3,doc4", - "docs": ["doc1", "doc2", "doc3", "doc4"], - } - }, ] @@ -3139,28 +2836,15 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: } assert [*app.stream({"query": "what is weather in sf"})] == [ + {"__start__": {"query": "what is weather in sf"}}, {"rewrite_query": {"query": "query: what is weather in sf"}}, { "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, "retriever_two": {"docs": ["doc3", "doc4"]}, "qa": {"answer": ""}, }, - { - "__end__": { - "answer": "", - "docs": ["doc3", "doc4"], - "query": "analyzed: query: what is weather in sf", - } - }, {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - { - "__end__": { - "query": "analyzed: query: what is weather in sf", - "answer": "doc1,doc2,doc3,doc4", - "docs": ["doc1", "doc2", "doc3", "doc4"], - } - }, ] app_w_interrupt = workflow.compile( @@ -3171,31 +2855,18 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: assert [ c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config) ] == [ + {"__start__": {"query": "what is weather in sf"}}, {"rewrite_query": {"query": "query: what is weather in sf"}}, { "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, "retriever_two": {"docs": ["doc3", "doc4"]}, "qa": {"answer": ""}, }, - { - "__end__": { - "answer": "", - "docs": ["doc3", "doc4"], - "query": "analyzed: query: what is weather in sf", - } - }, {"retriever_one": {"docs": ["doc1", "doc2"]}}, ] assert [c for c in app_w_interrupt.stream(None, config)] == [ {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - { - "__end__": { - "query": "analyzed: query: what is weather in sf", - "answer": "doc1,doc2,doc3,doc4", - "docs": ["doc1", "doc2", "doc3", "doc4"], - } - }, ] @@ -3264,13 +2935,13 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: } assert [*app.stream({"query": "what is weather in sf"})] == [ + {"__start__": {"query": "what is weather in sf"}}, {"rewrite_query": {"query": "query: what is weather in sf"}}, { "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, "retriever_two": {"docs": ["doc3", "doc4"]}, }, {"retriever_one": {"docs": ["doc1", "doc2"]}}, - {"decider": None}, {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}}, { "analyzer_one": { @@ -3281,22 +2952,5 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: { "retriever_one": {"docs": ["doc1", "doc2"]}, }, - {"decider": None}, {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}}, - { - "__end__": { - "query": "analyzed: query: analyzed: query: what is weather in sf", - "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4", - "docs": [ - "doc1", - "doc1", - "doc2", - "doc2", - "doc3", - "doc3", - "doc4", - "doc4", - ], - } - }, ] diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index e7b85ee04..c969f5261 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -47,8 +47,8 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: "input": LastValue(int), "output": LastValue(int), }, - input="input", - output="output", + input_channels="input", + output_channels="output", ) graph = Graph() graph.add_node("add_one", add_one) @@ -85,7 +85,9 @@ async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> N | Channel.write_to("output", fixed=5, output_plus_one=lambda x: x + 1) ) - app = Pregel(nodes={"one": chain}, output=["output", "fixed", "output_plus_one"]) + app = Pregel( + nodes={"one": chain}, output_channels=["output", "fixed", "output_plus_one"] + ) assert app.input_schema.schema() == {"title": "LangGraphInput"} assert app.output_schema.schema() == { @@ -128,7 +130,7 @@ async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: app = Pregel( nodes={"one": chain}, - output=["output"], + output_channels=["output"], ) assert app.input_schema.schema() == {"title": "LangGraphInput"} @@ -148,8 +150,8 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> nodes={ "one": chain, }, - input=["input"], - output=["output"], + input_channels=["input"], + output_channels=["output"], ) assert app.input_schema.schema() == { @@ -213,11 +215,7 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert values == { "add_one_more": 4, } - elif step == 3: - assert values == { - "__end__": 4, - } - assert step == 3 + assert step == 2 async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: @@ -281,17 +279,33 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: | Channel.write_to("output").abatch ) - pubsub = Pregel( + app = Pregel( nodes={"one": one, "two": two}, channels={"inbox": Topic(int)}, - input=["input", "inbox"], + input_channels=["input", "inbox"], ) # [12 + 1, 2 + 1 + 1] assert [ - c async for c in pubsub.astream({"input": 2, "inbox": 12}, output_keys="output") + c + async for c in app.astream( + {"input": 2, "inbox": 12}, output_keys="output", stream_mode="updates" + ) + ] == [ + {"two": 13}, + {"two": 4}, + ] + assert [ + c async for c in app.astream({"input": 2, "inbox": 12}, output_keys="output") ] == [13, 4] - assert [c async for c in pubsub.astream({"input": 2, "inbox": 12})] == [ + + assert [ + c async for c in app.astream({"input": 2, "inbox": 12}, stream_mode="updates") + ] == [ + {"one": {"inbox": 3}, "two": {"output": 13}}, + {"two": {"output": 4}}, + ] + assert [c async for c in app.astream({"input": 2, "inbox": 12})] == [ {"inbox": [3], "output": 13}, {"output": 4}, ] @@ -704,7 +718,7 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: "inbox": Topic(int), "ctx": Context(an_int, an_int_async, typ=int), }, - output=["inbox", "output"], + output_channels=["inbox", "output"], ) async def aenumerate(aiter: AsyncIterator[Any]) -> AsyncIterator[tuple[int, Any]]: @@ -928,32 +942,6 @@ async def test_conditional_graph() -> None: ), } }, - { - "__end__": { - "input": "what is weather in sf", - "intermediate_steps": [ - ( - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ), - ( - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ), - ], - "agent_outcome": AgentFinish( - return_values={"answer": "answer"}, log="finish:answer" - ), - } - }, ] patches = [c async for c in app.astream_log({"input": "what is weather in sf"})] @@ -1090,27 +1078,7 @@ async def test_conditional_graph() -> None: }, ) - 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", - ) - ], - } - } - ] + assert [c async for c in app_w_interrupt.astream(None, config)] == [] # test state get/update methods with interrupt_before @@ -1241,27 +1209,7 @@ async def test_conditional_graph() -> None: }, ) - 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", - ) - ], - } - } - ] + assert [c async for c in app_w_interrupt.astream(None, config)] == [] # test re-invoke to continue with interrupt_before @@ -1389,32 +1337,6 @@ async def test_conditional_graph() -> None: ), } }, - { - "__end__": { - "input": "what is weather in sf", - "intermediate_steps": [ - ( - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ), - ( - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ), - ], - "agent_outcome": AgentFinish( - return_values={"answer": "answer"}, log="finish:answer" - ), - } - }, ] @@ -1524,6 +1446,7 @@ async def test_conditional_graph_state() -> None: } assert [c async for c in app.astream({"input": "what is weather in sf"})] == [ + {"__start__": {"input": "what is weather in sf"}}, { "agent": { "agent_outcome": AgentAction( @@ -1575,32 +1498,6 @@ async def test_conditional_graph_state() -> None: ), } }, - { - "__end__": { - "input": "what is weather in sf", - "intermediate_steps": [ - ( - AgentAction( - tool="search_api", - tool_input="query", - log="tool:search_api:query", - ), - "result for query", - ), - ( - AgentAction( - tool="search_api", - tool_input="another", - log="tool:search_api:another", - ), - "result for another", - ), - ], - "agent_outcome": AgentFinish( - return_values={"answer": "answer"}, log="finish:answer" - ), - } - }, ] # test state get/update methods with interrupt_after @@ -1616,13 +1513,14 @@ async def test_conditional_graph_state() -> None: {"input": "what is weather in sf"}, config ) ] == [ + {"__start__": {"input": "what is weather in sf"}}, { "agent": { "agent_outcome": AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query" ), } - } + }, ] assert await app_w_interrupt.aget_state(config) == StateSnapshot( @@ -1698,27 +1596,7 @@ async def test_conditional_graph_state() -> None: }, ) - 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", - ) - ], - } - } - ] + assert [c async for c in app_w_interrupt.astream(None, config)] == [] # test state get/update methods with interrupt_before @@ -1734,13 +1612,14 @@ async def test_conditional_graph_state() -> None: {"input": "what is weather in sf"}, config ) ] == [ + {"__start__": {"input": "what is weather in sf"}}, { "agent": { "agent_outcome": AgentAction( tool="search_api", tool_input="query", log="tool:search_api:query" ), } - } + }, ] assert await app_w_interrupt.aget_state(config) == StateSnapshot( @@ -1816,27 +1695,7 @@ async def test_conditional_graph_state() -> None: }, ) - 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", - ) - ], - } - } - ] + assert [c async for c in app_w_interrupt.astream(None, config)] == [] async def test_conditional_entrypoint_graph() -> None: @@ -1872,7 +1731,6 @@ async def test_conditional_entrypoint_graph() -> None: assert [c async for c in app.astream("what is weather in sf")] == [ {"right": "what is weather in sf->right"}, - {"__end__": "what is weather in sf->right"}, ] @@ -1915,13 +1773,8 @@ async def test_conditional_entrypoint_graph_state() -> None: } assert [c async for c in app.astream({"input": "what is weather in sf"})] == [ + {"__start__": {"input": "what is weather in sf"}}, {"right": {"output": "what is weather in sf->right"}}, - { - "__end__": { - "input": "what is weather in sf", - "output": "what is weather in sf->right", - } - }, ] @@ -2044,6 +1897,7 @@ async def test_prebuilt_tool_chat() -> None: {"messages": [HumanMessage(content="what is weather in sf")]} ) ] == [ + {"__start__": {"messages": [HumanMessage(content="what is weather in sf")]}}, { "agent": { "messages": [ @@ -2114,61 +1968,6 @@ async def test_prebuilt_tool_chat() -> None: } }, {"agent": {"messages": [AIMessage(content="answer")]}}, - { - "__end__": { - "messages": [ - HumanMessage(content="what is weather in sf"), - AIMessage( - content="", - additional_kwargs={ - "tool_calls": [ - { - "id": "tool_call123", - "type": "function", - "function": { - "name": "search_api", - "arguments": '"query"', - }, - } - ] - }, - ), - ToolMessage( - content="result for query", tool_call_id="tool_call123" - ), - AIMessage( - content="", - additional_kwargs={ - "tool_calls": [ - { - "id": "tool_call234", - "type": "function", - "function": { - "name": "search_api", - "arguments": '"another"', - }, - }, - { - "id": "tool_call567", - "type": "function", - "function": { - "name": "search_api", - "arguments": '"a third one"', - }, - }, - ] - }, - ), - ToolMessage( - content="result for another", tool_call_id="tool_call234" - ), - ToolMessage( - content="result for a third one", tool_call_id="tool_call567" - ), - AIMessage(content="answer"), - ] - } - }, ] @@ -2244,6 +2043,7 @@ async def test_prebuilt_chat() -> None: {"messages": [HumanMessage(content="what is weather in sf")]} ) ] == [ + {"__start__": {"messages": [HumanMessage(content="what is weather in sf")]}}, { "agent": { "messages": [ @@ -2289,34 +2089,6 @@ async def test_prebuilt_chat() -> None: } }, {"agent": {"messages": [AIMessage(content="answer")]}}, - { - "__end__": { - "messages": [ - HumanMessage(content="what is weather in sf"), - AIMessage( - content="", - additional_kwargs={ - "function_call": { - "name": "search_api", - "arguments": '"query"', - } - }, - ), - FunctionMessage(content="result for query", name="search_api"), - AIMessage( - content="", - additional_kwargs={ - "function_call": { - "name": "search_api", - "arguments": '"another"', - } - }, - ), - FunctionMessage(content="result for another", name="search_api"), - AIMessage(content="answer"), - ] - } - }, ] @@ -2467,6 +2239,14 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None: assert [ c async for c in app.astream([HumanMessage(content="what is weather in sf")]) ] == [ + { + "__start__": [ + HumanMessage( + content="what is weather in sf", + id="00000000-0000-4000-8000-000000000038", + ) + ] + }, { "agent": AIMessage( content="", @@ -2500,42 +2280,6 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None: ) }, {"agent": AIMessage(content="answer", id="ai3")}, - { - "__end__": [ - HumanMessage( - content="what is weather in sf", - id="00000000-0000-4000-8000-000000000038", - ), - AIMessage( - content="", - additional_kwargs={ - "function_call": {"name": "search_api", "arguments": '"query"'} - }, - id="ai1", - ), - FunctionMessage( - content="result for query", - name="search_api", - id="00000000-0000-4000-8000-000000000051", - ), - AIMessage( - content="", - additional_kwargs={ - "function_call": { - "name": "search_api", - "arguments": '"another"', - } - }, - id="ai2", - ), - FunctionMessage( - content="result for another", - name="search_api", - id="00000000-0000-4000-8000-000000000064", - ), - AIMessage(content="answer", id="ai3"), - ] - }, ] app_w_interrupt = workflow.compile( @@ -2549,6 +2293,12 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None: HumanMessage(content="what is weather in sf"), config ) ] == [ + { + "__start__": HumanMessage( + content="what is weather in sf", + id="00000000-0000-4000-8000-000000000074", + ) + }, { "agent": AIMessage( content="", @@ -2557,7 +2307,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None: }, id="ai1", ) - } + }, ] assert await app_w_interrupt.aget_state(config) == StateSnapshot( @@ -2690,32 +2440,7 @@ async def test_message_graph(deterministic_uuids: MockerFixture) -> None: config=app_w_interrupt.checkpointer.get_tuple(config).config, ) - assert [c async for c in app_w_interrupt.astream(None, config)] == [ - { - "__end__": [ - HumanMessage( - content="what is weather in sf", - id="00000000-0000-4000-8000-000000000074", - ), - AIMessage( - content="", - additional_kwargs={ - "function_call": { - "name": "search_api", - "arguments": '"a different query"', - } - }, - id="ai1", - ), - FunctionMessage( - content="result for a different query", - name="search_api", - id="00000000-0000-4000-8000-000000000088", - ), - AIMessage(content="answer", id="ai2"), - ] - } - ] + assert [c async for c in app_w_interrupt.astream(None, config)] == [] async def test_in_one_fan_out_out_one_graph_state() -> None: @@ -2762,19 +2487,13 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: } assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ + {"__start__": {"query": "what is weather in sf"}}, {"rewrite_query": {"query": "query: what is weather in sf"}}, { "retriever_two": {"docs": ["doc3", "doc4"]}, "retriever_one": {"docs": ["doc1", "doc2"]}, }, {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - { - "__end__": { - "query": "query: what is weather in sf", - "answer": "doc1,doc2,doc3,doc4", - "docs": ["doc1", "doc2", "doc3", "doc4"], - } - }, ] @@ -2832,6 +2551,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge() -> None: } assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ + {"__start__": {"query": "what is weather in sf"}}, {"rewrite_query": {"query": "query: what is weather in sf"}}, { "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, @@ -2839,13 +2559,6 @@ async def test_in_one_fan_out_state_graph_waiting_edge() -> None: }, {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - { - "__end__": { - "query": "analyzed: query: what is weather in sf", - "answer": "doc1,doc2,doc3,doc4", - "docs": ["doc1", "doc2", "doc3", "doc4"], - } - }, ] app_w_interrupt = workflow.compile( @@ -2859,6 +2572,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge() -> None: {"query": "what is weather in sf"}, config ) ] == [ + {"__start__": {"query": "what is weather in sf"}}, {"rewrite_query": {"query": "query: what is weather in sf"}}, { "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, @@ -2869,13 +2583,6 @@ async def test_in_one_fan_out_state_graph_waiting_edge() -> None: assert [c async for c in app_w_interrupt.astream(None, config)] == [ {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - { - "__end__": { - "query": "analyzed: query: what is weather in sf", - "answer": "doc1,doc2,doc3,doc4", - "docs": ["doc1", "doc2", "doc3", "doc4"], - } - }, ] @@ -2937,28 +2644,15 @@ async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: } assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ + {"__start__": {"query": "what is weather in sf"}}, {"rewrite_query": {"query": "query: what is weather in sf"}}, { "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, "retriever_two": {"docs": ["doc3", "doc4"]}, "qa": {"answer": ""}, }, - { - "__end__": { - "answer": "", - "docs": ["doc3", "doc4"], - "query": "analyzed: query: what is weather in sf", - } - }, {"retriever_one": {"docs": ["doc1", "doc2"]}}, {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - { - "__end__": { - "query": "analyzed: query: what is weather in sf", - "answer": "doc1,doc2,doc3,doc4", - "docs": ["doc1", "doc2", "doc3", "doc4"], - } - }, ] app_w_interrupt = workflow.compile( @@ -2972,31 +2666,18 @@ async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular() -> None: {"query": "what is weather in sf"}, config ) ] == [ + {"__start__": {"query": "what is weather in sf"}}, {"rewrite_query": {"query": "query: what is weather in sf"}}, { "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, "retriever_two": {"docs": ["doc3", "doc4"]}, "qa": {"answer": ""}, }, - { - "__end__": { - "answer": "", - "docs": ["doc3", "doc4"], - "query": "analyzed: query: what is weather in sf", - } - }, {"retriever_one": {"docs": ["doc1", "doc2"]}}, ] assert [c async for c in app_w_interrupt.astream(None, config)] == [ {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, - { - "__end__": { - "query": "analyzed: query: what is weather in sf", - "answer": "doc1,doc2,doc3,doc4", - "docs": ["doc1", "doc2", "doc3", "doc4"], - } - }, ] @@ -3065,13 +2746,13 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: } assert [c async for c in app.astream({"query": "what is weather in sf"})] == [ + {"__start__": {"query": "what is weather in sf"}}, {"rewrite_query": {"query": "query: what is weather in sf"}}, { "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, "retriever_two": {"docs": ["doc3", "doc4"]}, }, {"retriever_one": {"docs": ["doc1", "doc2"]}}, - {"decider": None}, {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}}, { "analyzer_one": { @@ -3082,22 +2763,5 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: { "retriever_one": {"docs": ["doc1", "doc2"]}, }, - {"decider": None}, {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}}, - { - "__end__": { - "query": "analyzed: query: analyzed: query: what is weather in sf", - "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4", - "docs": [ - "doc1", - "doc1", - "doc2", - "doc2", - "doc3", - "doc3", - "doc4", - "doc4", - ], - } - }, ]