WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-13 05:07:51 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #245 from langchain-ai/nc/30mar/stream-mode

Browse Source 
Add stream_mode=updates or values
This commit is contained in:
Nuno Campos
2024-04-01 16:14:20 -07:00
committed by GitHub
parent 6de5a6de28 1ad017781d
commit e15ba64e1b
7 changed files with 288 additions and 909 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/graph/graph.py
+3 -4
View File
@@ -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,
langgraph/graph/state.py
+4 -7
View File
@@ -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,
langgraph/pregel/__init__.py
+86 -49
View File
@@ -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,
langgraph/pregel/io.py
+27 -1
View File
@@ -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
langgraph/pregel/validate.py
+25 -23
View File
@@ -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")
tests/test_pregel.py
+79 -425
View File
@@ -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",
],
}
},
]
tests/test_pregel_async.py
+64 -400
View File
@@ -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",
],
}
},
]