mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-20 06:35:46 +02:00
Rewrite graph drawing logic (#4354)
- It now executes the same pregel algo as when the graph is executed (without running any user code in nodes or conditional edges) to discover all the edges - This means we now support drawing the graph for all Pregel instances, not just StateGraph - This is done in preparation for new edge/node type coming in separate PR - Known changes - custom labels on conditional edges to END are no longer displayed
This commit is contained in:
Nuno Campos
@@ -1,3 +1,4 @@
|
||||
from collections.abc import Set
|
||||
from typing import Any, Generic, NamedTuple, Optional, Sequence, Type, Union
|
||||
|
||||
from typing_extensions import Self
|
||||
@@ -8,7 +9,7 @@ from langgraph.errors import EmptyChannelError, InvalidUpdateError
|
||||
|
||||
|
||||
class WaitForNames(NamedTuple):
|
||||
names: set[Any]
|
||||
names: Set[Any]
|
||||
|
||||
|
||||
class DynamicBarrierValue(
|
||||
@@ -25,7 +26,7 @@ class DynamicBarrierValue(
|
||||
|
||||
__slots__ = ("names", "seen")
|
||||
|
||||
names: Optional[set[Value]]
|
||||
names: Optional[Set[Value]]
|
||||
seen: set[Value]
|
||||
|
||||
def __init__(self, typ: Type[Value]) -> None:
|
||||
@@ -54,11 +55,11 @@ class DynamicBarrierValue(
|
||||
empty.seen = self.seen.copy()
|
||||
return empty
|
||||
|
||||
def checkpoint(self) -> tuple[Optional[set[Value]], set[Value]]:
|
||||
def checkpoint(self) -> tuple[Optional[Set[Value]], set[Value]]:
|
||||
return (self.names, self.seen)
|
||||
|
||||
def from_checkpoint(
|
||||
self, checkpoint: tuple[Optional[set[Value]], set[Value]]
|
||||
self, checkpoint: tuple[Optional[Set[Value]], set[Value]]
|
||||
) -> Self:
|
||||
empty = self.__class__(self.typ)
|
||||
empty.key = self.key
|
||||
|
||||
@@ -3,6 +3,7 @@ from inspect import (
|
||||
ismethod,
|
||||
signature,
|
||||
)
|
||||
from itertools import zip_longest
|
||||
from types import FunctionType
|
||||
from typing import (
|
||||
Any,
|
||||
@@ -29,12 +30,17 @@ from langchain_core.runnables import (
|
||||
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.errors import InvalidUpdateError
|
||||
from langgraph.pregel.write import ChannelWrite
|
||||
from langgraph.pregel.write import PASSTHROUGH, ChannelWrite, ChannelWriteEntry
|
||||
from langgraph.types import Send
|
||||
from langgraph.utils.runnable import (
|
||||
RunnableCallable,
|
||||
)
|
||||
|
||||
Writer = Callable[
|
||||
[Sequence[Union[str, Send]]],
|
||||
Sequence[Union[ChannelWriteEntry, Send]],
|
||||
]
|
||||
|
||||
|
||||
def _get_branch_path_input_schema(
|
||||
path: Union[
|
||||
@@ -124,9 +130,7 @@ class Branch(NamedTuple):
|
||||
|
||||
def run(
|
||||
self,
|
||||
writer: Callable[
|
||||
[Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
|
||||
],
|
||||
writer: Writer,
|
||||
reader: Optional[Callable[[RunnableConfig], Any]] = None,
|
||||
) -> RunnableCallable:
|
||||
return ChannelWrite.register_writer(
|
||||
@@ -138,7 +142,15 @@ class Branch(NamedTuple):
|
||||
name=None,
|
||||
trace=False,
|
||||
func_accepts_config=True,
|
||||
),
|
||||
list(
|
||||
zip_longest(
|
||||
writer([e for e in self.ends.values() if e != END]),
|
||||
[str(la) for la, e in self.ends.items() if e != END],
|
||||
)
|
||||
)
|
||||
if self.ends
|
||||
else None,
|
||||
)
|
||||
|
||||
def _route(
|
||||
@@ -147,9 +159,7 @@ class Branch(NamedTuple):
|
||||
config: RunnableConfig,
|
||||
*,
|
||||
reader: Optional[Callable[[RunnableConfig], Any]],
|
||||
writer: Callable[
|
||||
[Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
|
||||
],
|
||||
writer: Writer,
|
||||
) -> Runnable:
|
||||
if reader:
|
||||
value = reader(config)
|
||||
@@ -172,9 +182,7 @@ class Branch(NamedTuple):
|
||||
config: RunnableConfig,
|
||||
*,
|
||||
reader: Optional[Callable[[RunnableConfig], Any]],
|
||||
writer: Callable[
|
||||
[Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
|
||||
],
|
||||
writer: Writer,
|
||||
) -> Runnable:
|
||||
if reader:
|
||||
value = reader(config)
|
||||
@@ -193,9 +201,7 @@ class Branch(NamedTuple):
|
||||
|
||||
def _finish(
|
||||
self,
|
||||
writer: Callable[
|
||||
[Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
|
||||
],
|
||||
writer: Writer,
|
||||
input: Any,
|
||||
result: Any,
|
||||
config: RunnableConfig,
|
||||
@@ -212,4 +218,18 @@ class Branch(NamedTuple):
|
||||
raise ValueError("Branch did not return a valid destination")
|
||||
if any(p.node == END for p in destinations if isinstance(p, Send)):
|
||||
raise InvalidUpdateError("Cannot send a packet to the END node")
|
||||
return writer(destinations, config) or input
|
||||
entries = writer(destinations)
|
||||
if not entries:
|
||||
return input
|
||||
else:
|
||||
need_passthrough = False
|
||||
for e in entries:
|
||||
if isinstance(e, ChannelWriteEntry):
|
||||
if e.value is PASSTHROUGH:
|
||||
need_passthrough = True
|
||||
break
|
||||
if need_passthrough:
|
||||
return ChannelWrite(entries)
|
||||
else:
|
||||
ChannelWrite.do_write(config, entries)
|
||||
return input
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
import asyncio
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
from typing import (
|
||||
@@ -15,9 +14,6 @@ from typing import (
|
||||
)
|
||||
|
||||
from langchain_core.runnables import Runnable
|
||||
from langchain_core.runnables.config import RunnableConfig
|
||||
from langchain_core.runnables.graph import Graph as DrawableGraph
|
||||
from langchain_core.runnables.graph import Node as DrawableNode
|
||||
from typing_extensions import Self
|
||||
|
||||
from langgraph.channels.ephemeral_value import EphemeralValue
|
||||
@@ -32,7 +28,6 @@ from langgraph.constants import (
|
||||
)
|
||||
from langgraph.graph.branch import Branch
|
||||
from langgraph.pregel import Channel, Pregel
|
||||
from langgraph.pregel.protocol import PregelProtocol
|
||||
from langgraph.pregel.read import PregelNode
|
||||
from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
|
||||
from langgraph.types import All, Checkpointer
|
||||
@@ -380,10 +375,10 @@ class CompiledGraph(Pregel):
|
||||
cast(list[str], self.nodes[end].channels).append(start)
|
||||
|
||||
def attach_branch(self, start: str, name: str, branch: Branch) -> None:
|
||||
def branch_writer(
|
||||
packets: Sequence[Union[str, Send]], config: RunnableConfig
|
||||
) -> Optional[ChannelWrite]:
|
||||
writes = [
|
||||
def get_writes(
|
||||
packets: Sequence[Union[str, Send]],
|
||||
) -> Sequence[Union[ChannelWriteEntry, Send]]:
|
||||
return [
|
||||
(
|
||||
ChannelWriteEntry(f"branch:{start}:{name}:{p}" if p != END else END)
|
||||
if not isinstance(p, Send)
|
||||
@@ -391,14 +386,13 @@ class CompiledGraph(Pregel):
|
||||
)
|
||||
for p in packets
|
||||
]
|
||||
return ChannelWrite(cast(Sequence[Union[ChannelWriteEntry, Send]], writes))
|
||||
|
||||
# add hidden start node
|
||||
if start == START and start not in self.nodes:
|
||||
self.nodes[start] = Channel.subscribe_to(START, tags=[TAG_HIDDEN])
|
||||
|
||||
# attach branch writer
|
||||
self.nodes[start] |= branch.run(branch_writer)
|
||||
self.nodes[start] |= branch.run(get_writes)
|
||||
|
||||
# attach branch readers
|
||||
ends = branch.ends.values() if branch.ends else [node for node in self.nodes]
|
||||
@@ -408,171 +402,3 @@ class CompiledGraph(Pregel):
|
||||
self.channels[channel_name] = EphemeralValue(Any)
|
||||
self.nodes[end].triggers.append(channel_name)
|
||||
cast(list[str], self.nodes[end].channels).append(channel_name)
|
||||
|
||||
async def aget_graph(
|
||||
self,
|
||||
config: Optional[RunnableConfig] = None,
|
||||
*,
|
||||
xray: Union[int, bool] = False,
|
||||
) -> DrawableGraph:
|
||||
"""Returns a drawable representation of the computation graph."""
|
||||
from langgraph.pregel.remote import RemoteGraph
|
||||
|
||||
# gather subgraphs
|
||||
if xray:
|
||||
subpregels: dict[str, PregelProtocol] = {
|
||||
k: v
|
||||
async for k, v in self.aget_subgraphs()
|
||||
if isinstance(v, (CompiledGraph, RemoteGraph))
|
||||
}
|
||||
subgraphs = {
|
||||
k: v
|
||||
for k, v in zip(
|
||||
subpregels,
|
||||
await asyncio.gather(
|
||||
*(
|
||||
p.aget_graph(
|
||||
config,
|
||||
xray=xray
|
||||
if isinstance(xray, bool) or xray <= 0
|
||||
else xray - 1,
|
||||
)
|
||||
for p in subpregels.values()
|
||||
)
|
||||
),
|
||||
)
|
||||
}
|
||||
else:
|
||||
subgraphs = {}
|
||||
|
||||
# draw the graph
|
||||
return self._draw_graph(config, subgraphs=subgraphs)
|
||||
|
||||
def get_graph(
|
||||
self,
|
||||
config: Optional[RunnableConfig] = None,
|
||||
*,
|
||||
xray: Union[int, bool] = False,
|
||||
) -> DrawableGraph:
|
||||
"""Returns a drawable representation of the computation graph."""
|
||||
from langgraph.pregel.remote import RemoteGraph
|
||||
|
||||
# gather subgraphs
|
||||
if xray:
|
||||
subgraphs = {
|
||||
k: v.get_graph(
|
||||
config,
|
||||
xray=xray if isinstance(xray, bool) or xray <= 0 else xray - 1,
|
||||
)
|
||||
for k, v in self.get_subgraphs()
|
||||
if isinstance(v, (CompiledGraph, RemoteGraph))
|
||||
}
|
||||
else:
|
||||
subgraphs = {}
|
||||
|
||||
# draw the graph
|
||||
return self._draw_graph(config, subgraphs=subgraphs)
|
||||
|
||||
def _draw_graph(
|
||||
self,
|
||||
config: Optional[RunnableConfig] = None,
|
||||
*,
|
||||
subgraphs: dict[str, DrawableGraph] = {},
|
||||
) -> DrawableGraph:
|
||||
# create the graph
|
||||
graph = DrawableGraph()
|
||||
start_nodes: dict[str, DrawableNode] = {
|
||||
START: graph.add_node(self.get_input_schema(config), START)
|
||||
}
|
||||
end_nodes: dict[str, DrawableNode] = {}
|
||||
|
||||
def add_edge(
|
||||
start: str,
|
||||
end: str,
|
||||
label: Optional[Hashable] = None,
|
||||
conditional: bool = False,
|
||||
) -> None:
|
||||
if end == END and END not in end_nodes:
|
||||
end_nodes[END] = graph.add_node(self.get_output_schema(config), END)
|
||||
if start not in start_nodes or end not in end_nodes:
|
||||
logger.warning(
|
||||
f"Could not add edge from '{start}' to '{end}' due to missing nodes"
|
||||
)
|
||||
return
|
||||
return graph.add_edge(
|
||||
start_nodes[start],
|
||||
end_nodes[end],
|
||||
str(label) if label is not None else None,
|
||||
conditional,
|
||||
)
|
||||
|
||||
for key, n in self.builder.nodes.items():
|
||||
node = n.runnable
|
||||
metadata = n.metadata or {}
|
||||
if key in self.interrupt_before_nodes and key in self.interrupt_after_nodes:
|
||||
metadata["__interrupt"] = "before,after"
|
||||
elif key in self.interrupt_before_nodes:
|
||||
metadata["__interrupt"] = "before"
|
||||
elif key in self.interrupt_after_nodes:
|
||||
metadata["__interrupt"] = "after"
|
||||
if key in subgraphs:
|
||||
subgraph = subgraphs[key]
|
||||
subgraph.trim_first_node()
|
||||
subgraph.trim_last_node()
|
||||
if len(subgraph.nodes) >= 1:
|
||||
e, s = graph.extend(subgraph, prefix=key)
|
||||
if e is None:
|
||||
logger.warning(
|
||||
f"Could not extend subgraph '{key}' due to missing entrypoint"
|
||||
)
|
||||
continue
|
||||
if s is not None:
|
||||
start_nodes[key] = s
|
||||
end_nodes[key] = e
|
||||
else:
|
||||
nn = graph.add_node(node, key, metadata=metadata or None)
|
||||
start_nodes[key] = nn
|
||||
end_nodes[key] = nn
|
||||
else:
|
||||
nn = graph.add_node(node, key, metadata=metadata or None)
|
||||
start_nodes[key] = nn
|
||||
end_nodes[key] = nn
|
||||
for start, end in sorted(self.builder._all_edges):
|
||||
add_edge(start, end)
|
||||
for start, branches in self.builder.branches.items():
|
||||
default_ends = {
|
||||
**{k: k for k in self.builder.nodes if k != start},
|
||||
END: END,
|
||||
}
|
||||
for _, branch in branches.items():
|
||||
if branch.ends is not None:
|
||||
ends = branch.ends
|
||||
elif branch.then is not None:
|
||||
ends = {k: k for k in default_ends if k not in (END, branch.then)}
|
||||
else:
|
||||
ends = cast(dict[Hashable, str], default_ends)
|
||||
for label, end in ends.items():
|
||||
add_edge(
|
||||
start,
|
||||
end,
|
||||
label if label != end else None,
|
||||
conditional=True,
|
||||
)
|
||||
if branch.then is not None:
|
||||
add_edge(end, branch.then)
|
||||
for key, n in self.builder.nodes.items():
|
||||
if isinstance(n.ends, dict):
|
||||
for end, label in n.ends.items():
|
||||
add_edge(key, end, label, conditional=True)
|
||||
elif isinstance(n.ends, tuple):
|
||||
for end in n.ends:
|
||||
add_edge(key, end, conditional=True)
|
||||
|
||||
return graph
|
||||
|
||||
def _repr_mimebundle_(self, **kwargs: Any) -> dict[str, Any]:
|
||||
"""Mime bundle used by Jupyter to display the graph"""
|
||||
return {
|
||||
"text/plain": repr(self),
|
||||
"image/png": self.get_graph().draw_mermaid_png(),
|
||||
}
|
||||
|
||||
@@ -774,7 +774,12 @@ class CompiledStateGraph(CompiledGraph):
|
||||
ChannelWriteTupleEntry(
|
||||
mapper=_get_root if output_keys == ["__root__"] else _get_updates
|
||||
),
|
||||
ChannelWriteTupleEntry(mapper=_control_branch),
|
||||
ChannelWriteTupleEntry(
|
||||
mapper=_control_branch,
|
||||
static=_control_static(node.ends)
|
||||
if node is not None and node.ends is not None
|
||||
else None,
|
||||
),
|
||||
)
|
||||
|
||||
# add node and output channel
|
||||
@@ -840,9 +845,9 @@ class CompiledStateGraph(CompiledGraph):
|
||||
def attach_branch(
|
||||
self, start: str, name: str, branch: Branch, *, with_reader: bool = True
|
||||
) -> None:
|
||||
def branch_writer(
|
||||
packets: Sequence[Union[str, Send]], config: RunnableConfig
|
||||
) -> None:
|
||||
def get_writes(
|
||||
packets: Sequence[Union[str, Send]],
|
||||
) -> Sequence[Union[ChannelWriteEntry, Send]]:
|
||||
if filtered := [p for p in packets if p != END]:
|
||||
writes = [
|
||||
(
|
||||
@@ -857,13 +862,15 @@ class CompiledStateGraph(CompiledGraph):
|
||||
ChannelWriteEntry(
|
||||
f"branch:{start}:{name}::then",
|
||||
WaitForNames(
|
||||
{p.node if isinstance(p, Send) else p for p in filtered}
|
||||
frozenset(
|
||||
p.node if isinstance(p, Send) else p
|
||||
for p in filtered
|
||||
)
|
||||
),
|
||||
)
|
||||
)
|
||||
ChannelWrite.do_write(
|
||||
config, cast(Sequence[Union[Send, ChannelWriteEntry]], writes)
|
||||
)
|
||||
return writes
|
||||
return []
|
||||
|
||||
if with_reader:
|
||||
# get schema
|
||||
@@ -891,7 +898,7 @@ class CompiledStateGraph(CompiledGraph):
|
||||
reader = None
|
||||
|
||||
# attach branch publisher
|
||||
self.nodes[start].writers.append(branch.run(branch_writer, reader))
|
||||
self.nodes[start].writers.append(branch.run(get_writes, reader))
|
||||
|
||||
# attach then subscriber
|
||||
if branch.then and branch.then != END:
|
||||
@@ -1059,6 +1066,19 @@ def _control_branch(value: Any) -> Sequence[tuple[str, Any]]:
|
||||
return rtn
|
||||
|
||||
|
||||
def _control_static(
|
||||
ends: Union[tuple[str, ...], dict[str, str]],
|
||||
) -> Sequence[tuple[str, Any, Optional[str]]]:
|
||||
if isinstance(ends, dict):
|
||||
return [
|
||||
(CHANNEL_BRANCH_TO.format(k), None, label)
|
||||
for k, label in ends.items()
|
||||
if k != END
|
||||
]
|
||||
else:
|
||||
return [(CHANNEL_BRANCH_TO.format(e), None, None) for e in ends if e != END]
|
||||
|
||||
|
||||
def _get_root(input: Any) -> Optional[Sequence[tuple[str, Any]]]:
|
||||
if isinstance(input, Command):
|
||||
if input.graph == Command.PARENT:
|
||||
|
||||
@@ -93,6 +93,7 @@ from langgraph.pregel.algo import (
|
||||
)
|
||||
from langgraph.pregel.checkpoint import create_checkpoint, empty_checkpoint
|
||||
from langgraph.pregel.debug import tasks_w_writes
|
||||
from langgraph.pregel.draw import draw_graph
|
||||
from langgraph.pregel.io import map_input, read_channels
|
||||
from langgraph.pregel.loop import AsyncPregelLoop, StreamProtocol, SyncPregelLoop
|
||||
from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
|
||||
@@ -562,14 +563,87 @@ class Pregel(PregelProtocol):
|
||||
self.validate()
|
||||
|
||||
def get_graph(
|
||||
self, config: Optional[RunnableConfig] = None, *, xray: Union[int, bool] = False
|
||||
self,
|
||||
config: Optional[RunnableConfig] = None,
|
||||
*,
|
||||
xray: Union[int, bool] = False,
|
||||
) -> Graph:
|
||||
raise NotImplementedError
|
||||
"""Returns a drawable representation of the computation graph."""
|
||||
# gather subgraphs
|
||||
if xray:
|
||||
subgraphs = {
|
||||
k: v.get_graph(
|
||||
config,
|
||||
xray=xray if isinstance(xray, bool) or xray <= 0 else xray - 1,
|
||||
)
|
||||
for k, v in self.get_subgraphs()
|
||||
}
|
||||
else:
|
||||
subgraphs = {}
|
||||
|
||||
return draw_graph(
|
||||
merge_configs(self.config, config),
|
||||
nodes=self.nodes,
|
||||
specs=self.channels,
|
||||
input_channels=self.input_channels,
|
||||
interrupt_after_nodes=self.interrupt_after_nodes,
|
||||
interrupt_before_nodes=self.interrupt_before_nodes,
|
||||
trigger_to_nodes=self.trigger_to_nodes,
|
||||
checkpointer=self.checkpointer,
|
||||
subgraphs=subgraphs,
|
||||
)
|
||||
|
||||
async def aget_graph(
|
||||
self, config: Optional[RunnableConfig] = None, *, xray: Union[int, bool] = False
|
||||
self,
|
||||
config: Optional[RunnableConfig] = None,
|
||||
*,
|
||||
xray: Union[int, bool] = False,
|
||||
) -> Graph:
|
||||
raise NotImplementedError
|
||||
"""Returns a drawable representation of the computation graph."""
|
||||
|
||||
# gather subgraphs
|
||||
if xray:
|
||||
subpregels: dict[str, PregelProtocol] = {
|
||||
k: v async for k, v in self.aget_subgraphs()
|
||||
}
|
||||
subgraphs = {
|
||||
k: v
|
||||
for k, v in zip(
|
||||
subpregels,
|
||||
await asyncio.gather(
|
||||
*(
|
||||
p.aget_graph(
|
||||
config,
|
||||
xray=xray
|
||||
if isinstance(xray, bool) or xray <= 0
|
||||
else xray - 1,
|
||||
)
|
||||
for p in subpregels.values()
|
||||
)
|
||||
),
|
||||
)
|
||||
}
|
||||
else:
|
||||
subgraphs = {}
|
||||
|
||||
return draw_graph(
|
||||
merge_configs(self.config, config),
|
||||
nodes=self.nodes,
|
||||
specs=self.channels,
|
||||
input_channels=self.input_channels,
|
||||
interrupt_after_nodes=self.interrupt_after_nodes,
|
||||
interrupt_before_nodes=self.interrupt_before_nodes,
|
||||
trigger_to_nodes=self.trigger_to_nodes,
|
||||
checkpointer=self.checkpointer,
|
||||
subgraphs=subgraphs,
|
||||
)
|
||||
|
||||
def _repr_mimebundle_(self, **kwargs: Any) -> dict[str, Any]:
|
||||
"""Mime bundle used by Jupyter to display the graph"""
|
||||
return {
|
||||
"text/plain": repr(self),
|
||||
"image/png": self.get_graph().draw_mermaid_png(),
|
||||
}
|
||||
|
||||
def copy(self, update: Optional[dict[str, Any]] = None) -> Self:
|
||||
attrs = {**self.__dict__, **(update or {})}
|
||||
|
||||
@@ -0,0 +1,211 @@
|
||||
from collections import defaultdict
|
||||
from typing import Any, Mapping, Optional, Sequence, Union, cast
|
||||
|
||||
from langchain_core.runnables.config import RunnableConfig
|
||||
from langchain_core.runnables.graph import Graph, Node
|
||||
|
||||
from langgraph.channels.base import BaseChannel
|
||||
from langgraph.checkpoint.base import BaseCheckpointSaver
|
||||
from langgraph.constants import CONF, CONFIG_KEY_SEND, END, INPUT, START
|
||||
from langgraph.managed.base import ManagedValueSpec
|
||||
from langgraph.pregel.algo import (
|
||||
PregelTaskWrites,
|
||||
apply_writes,
|
||||
increment,
|
||||
prepare_next_tasks,
|
||||
)
|
||||
from langgraph.pregel.checkpoint import empty_checkpoint
|
||||
from langgraph.pregel.io import map_input
|
||||
from langgraph.pregel.manager import ChannelsManager
|
||||
from langgraph.pregel.read import PregelNode
|
||||
from langgraph.pregel.write import ChannelWrite
|
||||
from langgraph.types import All, Checkpointer, LoopProtocol
|
||||
|
||||
|
||||
def draw_graph(
|
||||
config: RunnableConfig,
|
||||
*,
|
||||
nodes: dict[str, PregelNode],
|
||||
specs: dict[str, Union[BaseChannel, ManagedValueSpec]],
|
||||
input_channels: Union[str, Sequence[str]],
|
||||
interrupt_after_nodes: Union[All, Sequence[str]],
|
||||
interrupt_before_nodes: Union[All, Sequence[str]],
|
||||
trigger_to_nodes: Optional[Mapping[str, Sequence[str]]],
|
||||
checkpointer: Checkpointer,
|
||||
subgraphs: dict[str, Graph],
|
||||
) -> Graph:
|
||||
"""Get the graph for this Pregel instance.
|
||||
|
||||
Args:
|
||||
config: The configuration to use for the graph.
|
||||
subgraphs: The subgraphs to include in the graph.
|
||||
checkpointer: The checkpointer to use for the graph.
|
||||
|
||||
Returns:
|
||||
The graph for this Pregel instance.
|
||||
"""
|
||||
# (src, dest, is_conditional, label)
|
||||
edges: set[tuple[str, str, bool, Optional[str]]] = set()
|
||||
|
||||
step = -1
|
||||
checkpoint = empty_checkpoint()
|
||||
get_next_version = (
|
||||
checkpointer.get_next_version
|
||||
if isinstance(checkpointer, BaseCheckpointSaver)
|
||||
else increment
|
||||
)
|
||||
with ChannelsManager(
|
||||
specs,
|
||||
checkpoint,
|
||||
LoopProtocol(step=step, stop=-1, config=config),
|
||||
skip_context=True,
|
||||
) as (channels, managed):
|
||||
static_seen: set[Any] = set()
|
||||
sources: dict[str, set[tuple[str, bool, Optional[str]]]] = {}
|
||||
step_sources: dict[str, set[tuple[str, bool, Optional[str]]]] = {}
|
||||
# remove node mappers
|
||||
nodes = {
|
||||
k: v.copy(update={"mapper": None}) if v.mapper is not None else v
|
||||
for k, v in nodes.items()
|
||||
}
|
||||
# apply input writes
|
||||
input_writes = list(map_input(input_channels, {}))
|
||||
_, updated_channels = apply_writes(
|
||||
checkpoint,
|
||||
channels,
|
||||
[
|
||||
PregelTaskWrites((), INPUT, input_writes, []),
|
||||
],
|
||||
get_next_version,
|
||||
)
|
||||
# prepare first tasks
|
||||
tasks = prepare_next_tasks(
|
||||
checkpoint,
|
||||
[],
|
||||
nodes,
|
||||
channels,
|
||||
managed,
|
||||
config,
|
||||
step,
|
||||
for_execution=True,
|
||||
store=None,
|
||||
checkpointer=None,
|
||||
manager=None,
|
||||
trigger_to_nodes=trigger_to_nodes,
|
||||
updated_channels=updated_channels,
|
||||
)
|
||||
start_tasks = tasks
|
||||
# run the pregel loop
|
||||
while tasks:
|
||||
conditionals: dict[tuple[str, str, Any], Optional[str]] = {}
|
||||
# run task writers
|
||||
for task in tasks.values():
|
||||
for w in task.writers:
|
||||
# apply regular writes
|
||||
if isinstance(w, ChannelWrite):
|
||||
w.invoke(None, task.config)
|
||||
# apply conditional writes declared for static analysis, only once
|
||||
if w not in static_seen:
|
||||
static_seen.add(w)
|
||||
# apply static writes
|
||||
if writes := ChannelWrite.get_static_writes(w):
|
||||
conditionals.update(
|
||||
{(task.name, *t[:2]): t[2] for t in writes}
|
||||
)
|
||||
task.config[CONF][CONFIG_KEY_SEND]([t[:2] for t in writes])
|
||||
# collect sources
|
||||
step_sources = {
|
||||
task.name: {
|
||||
(
|
||||
w[0],
|
||||
(task.name, *w) in conditionals,
|
||||
conditionals.get((task.name, *w)),
|
||||
)
|
||||
for w in task.writes
|
||||
}
|
||||
for task in tasks.values()
|
||||
}
|
||||
sources.update(step_sources)
|
||||
# invert triggers
|
||||
trigger_to_sources: dict[str, set[tuple[str, bool, Optional[str]]]] = (
|
||||
defaultdict(set)
|
||||
)
|
||||
for src, triggers in sources.items():
|
||||
for trigger, cond, label in triggers:
|
||||
trigger_to_sources[trigger].add((src, cond, label))
|
||||
# apply writes
|
||||
_, updated_channels = apply_writes(
|
||||
checkpoint, channels, tasks.values(), get_next_version
|
||||
)
|
||||
# prepare next tasks
|
||||
tasks = prepare_next_tasks(
|
||||
checkpoint,
|
||||
[],
|
||||
nodes,
|
||||
channels,
|
||||
managed,
|
||||
config,
|
||||
step,
|
||||
for_execution=True,
|
||||
store=None,
|
||||
checkpointer=None,
|
||||
manager=None,
|
||||
trigger_to_nodes=trigger_to_nodes,
|
||||
updated_channels=updated_channels,
|
||||
)
|
||||
# collect edges
|
||||
for task in tasks.values():
|
||||
for trigger in task.triggers:
|
||||
for src, cond, label in sorted(trigger_to_sources[trigger]):
|
||||
edges.add((src, task.name, cond, label))
|
||||
# assemble the graph
|
||||
graph = Graph()
|
||||
# add nodes
|
||||
for name, node in nodes.items():
|
||||
metadata = dict(node.metadata or {})
|
||||
if name in interrupt_before_nodes and name in interrupt_after_nodes:
|
||||
metadata["__interrupt"] = "before,after"
|
||||
elif name in interrupt_before_nodes:
|
||||
metadata["__interrupt"] = "before"
|
||||
elif name in interrupt_after_nodes:
|
||||
metadata["__interrupt"] = "after"
|
||||
graph.add_node(node.bound, name, metadata=metadata or None)
|
||||
# add start node
|
||||
if START not in nodes:
|
||||
graph.add_node(None, START)
|
||||
for task in start_tasks.values():
|
||||
graph.add_edge(graph.nodes[START], graph.nodes[task.name])
|
||||
# add discovered edges
|
||||
for src, dest, is_conditional, label in sorted(edges):
|
||||
graph.add_edge(
|
||||
graph.nodes[src],
|
||||
graph.nodes[dest],
|
||||
data=label if label != dest else None,
|
||||
conditional=is_conditional,
|
||||
)
|
||||
# add end edges
|
||||
if step_sources:
|
||||
end = graph.add_node(None, END)
|
||||
termini = {d for _, d, _, _ in edges}.difference(s for s, _, _, _ in edges)
|
||||
for src in sorted(termini.union(step_sources)):
|
||||
graph.add_edge(graph.nodes[src], end, conditional=src not in termini)
|
||||
# replace subgraphs
|
||||
for name, subgraph in subgraphs.items():
|
||||
subgraph.trim_first_node()
|
||||
subgraph.trim_last_node()
|
||||
if (
|
||||
len(subgraph.nodes) > 1
|
||||
and name in graph.nodes
|
||||
and subgraph.first_node()
|
||||
and subgraph.last_node()
|
||||
):
|
||||
# replace the node with the subgraph
|
||||
graph.nodes.pop(name)
|
||||
first, last = graph.extend(subgraph, prefix=name)
|
||||
for idx, edge in enumerate(graph.edges):
|
||||
if edge.source == name:
|
||||
graph.edges[idx] = edge.copy(source=cast(Node, last).id)
|
||||
elif edge.target == name:
|
||||
graph.edges[idx] = edge.copy(target=cast(Node, first).id)
|
||||
|
||||
return graph
|
||||
@@ -14,7 +14,7 @@ from typing import (
|
||||
from langchain_core.runnables import Runnable, RunnableConfig
|
||||
from langchain_core.runnables.utils import ConfigurableFieldSpec
|
||||
|
||||
from langgraph.constants import CONF, CONFIG_KEY_SEND, TASKS, Send
|
||||
from langgraph.constants import CONF, CONFIG_KEY_SEND, MISSING, TASKS, Send
|
||||
from langgraph.errors import InvalidUpdateError
|
||||
from langgraph.utils.runnable import RunnableCallable
|
||||
|
||||
@@ -41,6 +41,8 @@ class ChannelWriteTupleEntry(NamedTuple):
|
||||
"""Function to extract tuples from value."""
|
||||
value: Any = PASSTHROUGH
|
||||
"""Value to write, or PASSTHROUGH to use the input."""
|
||||
static: Optional[Sequence[tuple[str, Any, Optional[str]]]] = None
|
||||
"""Optional, declared writes for static analysis."""
|
||||
|
||||
|
||||
class ChannelWrite(RunnableCallable):
|
||||
@@ -121,6 +123,7 @@ class ChannelWrite(RunnableCallable):
|
||||
def do_write(
|
||||
config: RunnableConfig,
|
||||
writes: Sequence[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]],
|
||||
allow_passthrough: bool = True,
|
||||
require_at_least_one_of: Optional[Sequence[str]] = None, # ignored
|
||||
) -> None:
|
||||
# validate
|
||||
@@ -130,46 +133,80 @@ class ChannelWrite(RunnableCallable):
|
||||
raise InvalidUpdateError(
|
||||
"Cannot write to the reserved channel TASKS"
|
||||
)
|
||||
if w.value is PASSTHROUGH:
|
||||
if w.value is PASSTHROUGH and not allow_passthrough:
|
||||
raise InvalidUpdateError("PASSTHROUGH value must be replaced")
|
||||
if isinstance(w, ChannelWriteTupleEntry):
|
||||
if w.value is PASSTHROUGH:
|
||||
if w.value is PASSTHROUGH and not allow_passthrough:
|
||||
raise InvalidUpdateError("PASSTHROUGH value must be replaced")
|
||||
# assemble writes
|
||||
tuples: list[tuple[str, Any]] = []
|
||||
for w in writes:
|
||||
if isinstance(w, Send):
|
||||
tuples.append((TASKS, w))
|
||||
elif isinstance(w, ChannelWriteTupleEntry):
|
||||
if ww := w.mapper(w.value):
|
||||
tuples.extend(ww)
|
||||
elif isinstance(w, ChannelWriteEntry):
|
||||
value = w.mapper(w.value) if w.mapper is not None else w.value
|
||||
if value is SKIP_WRITE:
|
||||
continue
|
||||
if w.skip_none and value is None:
|
||||
continue
|
||||
tuples.append((w.channel, value))
|
||||
else:
|
||||
raise ValueError(f"Invalid write entry: {w}")
|
||||
# if we want to persist writes found before hitting a ParentCommand
|
||||
# can move this to a finally block
|
||||
write: TYPE_SEND = config[CONF][CONFIG_KEY_SEND]
|
||||
write(tuples)
|
||||
write(_assemble_writes(writes))
|
||||
|
||||
@staticmethod
|
||||
def is_writer(runnable: Runnable) -> bool:
|
||||
"""Used by PregelNode to distinguish between writers and other runnables."""
|
||||
return (
|
||||
isinstance(runnable, ChannelWrite)
|
||||
or getattr(runnable, "_is_channel_writer", False) is True
|
||||
or getattr(runnable, "_is_channel_writer", MISSING) is not MISSING
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def register_writer(runnable: R) -> R:
|
||||
def get_static_writes(
|
||||
runnable: Runnable,
|
||||
) -> Optional[Sequence[tuple[str, Any, Optional[str]]]]:
|
||||
"""Used to get conditional writes a writer declares for static analysis."""
|
||||
if isinstance(runnable, ChannelWrite):
|
||||
return [
|
||||
w
|
||||
for entry in runnable.writes
|
||||
if isinstance(entry, ChannelWriteTupleEntry) and entry.static
|
||||
for w in entry.static
|
||||
] or None
|
||||
elif writes := getattr(runnable, "_is_channel_writer", MISSING):
|
||||
if writes is not MISSING:
|
||||
writes = cast(
|
||||
Sequence[tuple[Union[ChannelWriteEntry, Send], Optional[str]]],
|
||||
writes,
|
||||
)
|
||||
entries = [e for e, _ in writes]
|
||||
labels = [la for _, la in writes]
|
||||
return [(*t, la) for t, la in zip(_assemble_writes(entries), labels)]
|
||||
|
||||
@staticmethod
|
||||
def register_writer(
|
||||
runnable: R,
|
||||
static: Optional[
|
||||
Sequence[tuple[Union[ChannelWriteEntry, Send], Optional[str]]]
|
||||
] = None,
|
||||
) -> R:
|
||||
"""Used to mark a runnable as a writer, so that it can be detected by is_writer.
|
||||
Instances of ChannelWrite are automatically marked as writers."""
|
||||
Instances of ChannelWrite are automatically marked as writers.
|
||||
Optionally, a list of declared writes can be passed for static analysis."""
|
||||
# using object.__setattr__ to work around objects that override __setattr__
|
||||
# eg. pydantic models and dataclasses
|
||||
object.__setattr__(runnable, "_is_channel_writer", True)
|
||||
object.__setattr__(runnable, "_is_channel_writer", static)
|
||||
return runnable
|
||||
|
||||
|
||||
def _assemble_writes(
|
||||
writes: Sequence[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]],
|
||||
) -> list[tuple[str, Any]]:
|
||||
"""Assembles the writes into a list of tuples."""
|
||||
tuples: list[tuple[str, Any]] = []
|
||||
for w in writes:
|
||||
if isinstance(w, Send):
|
||||
tuples.append((TASKS, w))
|
||||
elif isinstance(w, ChannelWriteTupleEntry):
|
||||
if ww := w.mapper(w.value):
|
||||
tuples.extend(ww)
|
||||
elif isinstance(w, ChannelWriteEntry):
|
||||
value = w.mapper(w.value) if w.mapper is not None else w.value
|
||||
if value is SKIP_WRITE:
|
||||
continue
|
||||
if w.skip_none and value is None:
|
||||
continue
|
||||
tuples.append((w.channel, value))
|
||||
else:
|
||||
raise ValueError(f"Invalid write entry: {w}")
|
||||
return tuples
|
||||
|
||||
Generated
+4
-4
@@ -1,4 +1,4 @@
|
||||
# This file is automatically @generated by Poetry 2.0.1 and should not be changed by hand.
|
||||
# This file is automatically @generated by Poetry 2.0.0 and should not be changed by hand.
|
||||
|
||||
[[package]]
|
||||
name = "aiosqlite"
|
||||
@@ -1324,14 +1324,14 @@ files = [
|
||||
|
||||
[[package]]
|
||||
name = "langchain-core"
|
||||
version = "0.3.46"
|
||||
version = "0.3.55"
|
||||
description = "Building applications with LLMs through composability"
|
||||
optional = false
|
||||
python-versions = "<4.0,>=3.9"
|
||||
groups = ["main", "dev"]
|
||||
files = [
|
||||
{file = "langchain_core-0.3.46-py3-none-any.whl", hash = "sha256:28b5689fc347975ea520b5364ab4aee5567e661553bbee5e97cabf4596c28ce0"},
|
||||
{file = "langchain_core-0.3.46.tar.gz", hash = "sha256:5fca010eeb0a427be5aa8a8525e2112995dde790c584cef165be7c5e0ee1c2b5"},
|
||||
{file = "langchain_core-0.3.55-py3-none-any.whl", hash = "sha256:b3cb36bf37755a616158a79866657c6697b43a2f7c69dd723ce425f1c76c1baa"},
|
||||
{file = "langchain_core-0.3.55.tar.gz", hash = "sha256:0f2b3e311621116a83510c70b0ac9d959030a0a457a69483535cff18501fedc9"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -1,151 +0,0 @@
|
||||
# serializer version: 1
|
||||
# name: test_weather_subgraph[memory]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
router_node(router_node)
|
||||
normal_llm_node(normal_llm_node)
|
||||
__end__([<p>__end__</p>]):::last
|
||||
__start__ --> router_node;
|
||||
normal_llm_node --> __end__;
|
||||
weather_graph_weather_node --> __end__;
|
||||
router_node -.-> normal_llm_node;
|
||||
router_node -.-> weather_graph_model_node;
|
||||
router_node -.-> __end__;
|
||||
subgraph weather_graph
|
||||
weather_graph_model_node(model_node)
|
||||
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
|
||||
weather_graph_model_node --> weather_graph_weather_node;
|
||||
end
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_weather_subgraph[postgres_aio]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
router_node(router_node)
|
||||
normal_llm_node(normal_llm_node)
|
||||
__end__([<p>__end__</p>]):::last
|
||||
__start__ --> router_node;
|
||||
normal_llm_node --> __end__;
|
||||
weather_graph_weather_node --> __end__;
|
||||
router_node -.-> normal_llm_node;
|
||||
router_node -.-> weather_graph_model_node;
|
||||
router_node -.-> __end__;
|
||||
subgraph weather_graph
|
||||
weather_graph_model_node(model_node)
|
||||
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
|
||||
weather_graph_model_node --> weather_graph_weather_node;
|
||||
end
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_weather_subgraph[postgres_aio_pipe]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
router_node(router_node)
|
||||
normal_llm_node(normal_llm_node)
|
||||
__end__([<p>__end__</p>]):::last
|
||||
__start__ --> router_node;
|
||||
normal_llm_node --> __end__;
|
||||
weather_graph_weather_node --> __end__;
|
||||
router_node -.-> normal_llm_node;
|
||||
router_node -.-> weather_graph_model_node;
|
||||
router_node -.-> __end__;
|
||||
subgraph weather_graph
|
||||
weather_graph_model_node(model_node)
|
||||
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
|
||||
weather_graph_model_node --> weather_graph_weather_node;
|
||||
end
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_weather_subgraph[postgres_aio_pool]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
router_node(router_node)
|
||||
normal_llm_node(normal_llm_node)
|
||||
__end__([<p>__end__</p>]):::last
|
||||
__start__ --> router_node;
|
||||
normal_llm_node --> __end__;
|
||||
weather_graph_weather_node --> __end__;
|
||||
router_node -.-> normal_llm_node;
|
||||
router_node -.-> weather_graph_model_node;
|
||||
router_node -.-> __end__;
|
||||
subgraph weather_graph
|
||||
weather_graph_model_node(model_node)
|
||||
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
|
||||
weather_graph_model_node --> weather_graph_weather_node;
|
||||
end
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_weather_subgraph[postgres_aio_shallow]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
router_node(router_node)
|
||||
normal_llm_node(normal_llm_node)
|
||||
__end__([<p>__end__</p>]):::last
|
||||
__start__ --> router_node;
|
||||
normal_llm_node --> __end__;
|
||||
weather_graph_weather_node --> __end__;
|
||||
router_node -.-> normal_llm_node;
|
||||
router_node -.-> weather_graph_model_node;
|
||||
router_node -.-> __end__;
|
||||
subgraph weather_graph
|
||||
weather_graph_model_node(model_node)
|
||||
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
|
||||
weather_graph_model_node --> weather_graph_weather_node;
|
||||
end
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_weather_subgraph[sqlite_aio]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
router_node(router_node)
|
||||
normal_llm_node(normal_llm_node)
|
||||
__end__([<p>__end__</p>]):::last
|
||||
__start__ --> router_node;
|
||||
normal_llm_node --> __end__;
|
||||
weather_graph_weather_node --> __end__;
|
||||
router_node -.-> normal_llm_node;
|
||||
router_node -.-> weather_graph_model_node;
|
||||
router_node -.-> __end__;
|
||||
subgraph weather_graph
|
||||
weather_graph_model_node(model_node)
|
||||
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
|
||||
weather_graph_model_node --> weather_graph_weather_node;
|
||||
end
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
File diff suppressed because it is too large
Load Diff
@@ -4,11 +4,11 @@
|
||||
graph TD;
|
||||
__start__ --> rewrite_query;
|
||||
analyzer_one --> retriever_one;
|
||||
qa --> __end__;
|
||||
retriever_one --> qa;
|
||||
retriever_two --> qa;
|
||||
rewrite_query --> analyzer_one;
|
||||
rewrite_query -.-> retriever_two;
|
||||
qa --> __end__;
|
||||
|
||||
'''
|
||||
# ---
|
||||
@@ -120,695 +120,21 @@
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> rewrite_query;
|
||||
analyzer_one --> retriever_one;
|
||||
qa --> __end__;
|
||||
retriever_one --> qa;
|
||||
retriever_two --> qa;
|
||||
rewrite_query --> analyzer_one;
|
||||
rewrite_query -.-> retriever_two;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio].1
|
||||
dict({
|
||||
'$defs': dict({
|
||||
'InnerObject': dict({
|
||||
'properties': dict({
|
||||
'yo': dict({
|
||||
'title': 'Yo',
|
||||
'type': 'integer',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'yo',
|
||||
]),
|
||||
'title': 'InnerObject',
|
||||
'type': 'object',
|
||||
}),
|
||||
}),
|
||||
'properties': dict({
|
||||
'answer': dict({
|
||||
'anyOf': list([
|
||||
dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
dict({
|
||||
'type': 'null',
|
||||
}),
|
||||
]),
|
||||
'default': None,
|
||||
'title': 'Answer',
|
||||
}),
|
||||
'docs': dict({
|
||||
'items': dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
'title': 'Docs',
|
||||
'type': 'array',
|
||||
}),
|
||||
'inner': dict({
|
||||
'$ref': '#/$defs/InnerObject',
|
||||
}),
|
||||
'query': dict({
|
||||
'title': 'Query',
|
||||
'type': 'string',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'query',
|
||||
'inner',
|
||||
'docs',
|
||||
]),
|
||||
'title': 'State',
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio].2
|
||||
dict({
|
||||
'$defs': dict({
|
||||
'InnerObject': dict({
|
||||
'properties': dict({
|
||||
'yo': dict({
|
||||
'title': 'Yo',
|
||||
'type': 'integer',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'yo',
|
||||
]),
|
||||
'title': 'InnerObject',
|
||||
'type': 'object',
|
||||
}),
|
||||
}),
|
||||
'properties': dict({
|
||||
'answer': dict({
|
||||
'anyOf': list([
|
||||
dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
dict({
|
||||
'type': 'null',
|
||||
}),
|
||||
]),
|
||||
'default': None,
|
||||
'title': 'Answer',
|
||||
}),
|
||||
'docs': dict({
|
||||
'items': dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
'title': 'Docs',
|
||||
'type': 'array',
|
||||
}),
|
||||
'inner': dict({
|
||||
'$ref': '#/$defs/InnerObject',
|
||||
}),
|
||||
'query': dict({
|
||||
'title': 'Query',
|
||||
'type': 'string',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'query',
|
||||
'inner',
|
||||
'docs',
|
||||
]),
|
||||
'title': 'State',
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pipe]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> rewrite_query;
|
||||
analyzer_one --> retriever_one;
|
||||
qa --> __end__;
|
||||
retriever_one --> qa;
|
||||
retriever_two --> qa;
|
||||
rewrite_query --> analyzer_one;
|
||||
rewrite_query -.-> retriever_two;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pipe].1
|
||||
dict({
|
||||
'$defs': dict({
|
||||
'InnerObject': dict({
|
||||
'properties': dict({
|
||||
'yo': dict({
|
||||
'title': 'Yo',
|
||||
'type': 'integer',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'yo',
|
||||
]),
|
||||
'title': 'InnerObject',
|
||||
'type': 'object',
|
||||
}),
|
||||
}),
|
||||
'properties': dict({
|
||||
'answer': dict({
|
||||
'anyOf': list([
|
||||
dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
dict({
|
||||
'type': 'null',
|
||||
}),
|
||||
]),
|
||||
'default': None,
|
||||
'title': 'Answer',
|
||||
}),
|
||||
'docs': dict({
|
||||
'items': dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
'title': 'Docs',
|
||||
'type': 'array',
|
||||
}),
|
||||
'inner': dict({
|
||||
'$ref': '#/$defs/InnerObject',
|
||||
}),
|
||||
'query': dict({
|
||||
'title': 'Query',
|
||||
'type': 'string',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'query',
|
||||
'inner',
|
||||
'docs',
|
||||
]),
|
||||
'title': 'State',
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pipe].2
|
||||
dict({
|
||||
'$defs': dict({
|
||||
'InnerObject': dict({
|
||||
'properties': dict({
|
||||
'yo': dict({
|
||||
'title': 'Yo',
|
||||
'type': 'integer',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'yo',
|
||||
]),
|
||||
'title': 'InnerObject',
|
||||
'type': 'object',
|
||||
}),
|
||||
}),
|
||||
'properties': dict({
|
||||
'answer': dict({
|
||||
'anyOf': list([
|
||||
dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
dict({
|
||||
'type': 'null',
|
||||
}),
|
||||
]),
|
||||
'default': None,
|
||||
'title': 'Answer',
|
||||
}),
|
||||
'docs': dict({
|
||||
'items': dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
'title': 'Docs',
|
||||
'type': 'array',
|
||||
}),
|
||||
'inner': dict({
|
||||
'$ref': '#/$defs/InnerObject',
|
||||
}),
|
||||
'query': dict({
|
||||
'title': 'Query',
|
||||
'type': 'string',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'query',
|
||||
'inner',
|
||||
'docs',
|
||||
]),
|
||||
'title': 'State',
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pool]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> rewrite_query;
|
||||
analyzer_one --> retriever_one;
|
||||
qa --> __end__;
|
||||
retriever_one --> qa;
|
||||
retriever_two --> qa;
|
||||
rewrite_query --> analyzer_one;
|
||||
rewrite_query -.-> retriever_two;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pool].1
|
||||
dict({
|
||||
'$defs': dict({
|
||||
'InnerObject': dict({
|
||||
'properties': dict({
|
||||
'yo': dict({
|
||||
'title': 'Yo',
|
||||
'type': 'integer',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'yo',
|
||||
]),
|
||||
'title': 'InnerObject',
|
||||
'type': 'object',
|
||||
}),
|
||||
}),
|
||||
'properties': dict({
|
||||
'answer': dict({
|
||||
'anyOf': list([
|
||||
dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
dict({
|
||||
'type': 'null',
|
||||
}),
|
||||
]),
|
||||
'default': None,
|
||||
'title': 'Answer',
|
||||
}),
|
||||
'docs': dict({
|
||||
'items': dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
'title': 'Docs',
|
||||
'type': 'array',
|
||||
}),
|
||||
'inner': dict({
|
||||
'$ref': '#/$defs/InnerObject',
|
||||
}),
|
||||
'query': dict({
|
||||
'title': 'Query',
|
||||
'type': 'string',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'query',
|
||||
'inner',
|
||||
'docs',
|
||||
]),
|
||||
'title': 'State',
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_pool].2
|
||||
dict({
|
||||
'$defs': dict({
|
||||
'InnerObject': dict({
|
||||
'properties': dict({
|
||||
'yo': dict({
|
||||
'title': 'Yo',
|
||||
'type': 'integer',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'yo',
|
||||
]),
|
||||
'title': 'InnerObject',
|
||||
'type': 'object',
|
||||
}),
|
||||
}),
|
||||
'properties': dict({
|
||||
'answer': dict({
|
||||
'anyOf': list([
|
||||
dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
dict({
|
||||
'type': 'null',
|
||||
}),
|
||||
]),
|
||||
'default': None,
|
||||
'title': 'Answer',
|
||||
}),
|
||||
'docs': dict({
|
||||
'items': dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
'title': 'Docs',
|
||||
'type': 'array',
|
||||
}),
|
||||
'inner': dict({
|
||||
'$ref': '#/$defs/InnerObject',
|
||||
}),
|
||||
'query': dict({
|
||||
'title': 'Query',
|
||||
'type': 'string',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'query',
|
||||
'inner',
|
||||
'docs',
|
||||
]),
|
||||
'title': 'State',
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_shallow]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> rewrite_query;
|
||||
analyzer_one --> retriever_one;
|
||||
qa --> __end__;
|
||||
retriever_one --> qa;
|
||||
retriever_two --> qa;
|
||||
rewrite_query --> analyzer_one;
|
||||
rewrite_query -.-> retriever_two;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_shallow].1
|
||||
dict({
|
||||
'$defs': dict({
|
||||
'InnerObject': dict({
|
||||
'properties': dict({
|
||||
'yo': dict({
|
||||
'title': 'Yo',
|
||||
'type': 'integer',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'yo',
|
||||
]),
|
||||
'title': 'InnerObject',
|
||||
'type': 'object',
|
||||
}),
|
||||
}),
|
||||
'properties': dict({
|
||||
'answer': dict({
|
||||
'anyOf': list([
|
||||
dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
dict({
|
||||
'type': 'null',
|
||||
}),
|
||||
]),
|
||||
'default': None,
|
||||
'title': 'Answer',
|
||||
}),
|
||||
'docs': dict({
|
||||
'items': dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
'title': 'Docs',
|
||||
'type': 'array',
|
||||
}),
|
||||
'inner': dict({
|
||||
'$ref': '#/$defs/InnerObject',
|
||||
}),
|
||||
'query': dict({
|
||||
'title': 'Query',
|
||||
'type': 'string',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'query',
|
||||
'inner',
|
||||
'docs',
|
||||
]),
|
||||
'title': 'State',
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_aio_shallow].2
|
||||
dict({
|
||||
'$defs': dict({
|
||||
'InnerObject': dict({
|
||||
'properties': dict({
|
||||
'yo': dict({
|
||||
'title': 'Yo',
|
||||
'type': 'integer',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'yo',
|
||||
]),
|
||||
'title': 'InnerObject',
|
||||
'type': 'object',
|
||||
}),
|
||||
}),
|
||||
'properties': dict({
|
||||
'answer': dict({
|
||||
'anyOf': list([
|
||||
dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
dict({
|
||||
'type': 'null',
|
||||
}),
|
||||
]),
|
||||
'default': None,
|
||||
'title': 'Answer',
|
||||
}),
|
||||
'docs': dict({
|
||||
'items': dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
'title': 'Docs',
|
||||
'type': 'array',
|
||||
}),
|
||||
'inner': dict({
|
||||
'$ref': '#/$defs/InnerObject',
|
||||
}),
|
||||
'query': dict({
|
||||
'title': 'Query',
|
||||
'type': 'string',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'query',
|
||||
'inner',
|
||||
'docs',
|
||||
]),
|
||||
'title': 'State',
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite_aio]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> rewrite_query;
|
||||
analyzer_one --> retriever_one;
|
||||
qa --> __end__;
|
||||
retriever_one --> qa;
|
||||
retriever_two --> qa;
|
||||
rewrite_query --> analyzer_one;
|
||||
rewrite_query -.-> retriever_two;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite_aio].1
|
||||
dict({
|
||||
'$defs': dict({
|
||||
'InnerObject': dict({
|
||||
'properties': dict({
|
||||
'yo': dict({
|
||||
'title': 'Yo',
|
||||
'type': 'integer',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'yo',
|
||||
]),
|
||||
'title': 'InnerObject',
|
||||
'type': 'object',
|
||||
}),
|
||||
}),
|
||||
'properties': dict({
|
||||
'answer': dict({
|
||||
'anyOf': list([
|
||||
dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
dict({
|
||||
'type': 'null',
|
||||
}),
|
||||
]),
|
||||
'default': None,
|
||||
'title': 'Answer',
|
||||
}),
|
||||
'docs': dict({
|
||||
'items': dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
'title': 'Docs',
|
||||
'type': 'array',
|
||||
}),
|
||||
'inner': dict({
|
||||
'$ref': '#/$defs/InnerObject',
|
||||
}),
|
||||
'query': dict({
|
||||
'title': 'Query',
|
||||
'type': 'string',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'query',
|
||||
'inner',
|
||||
'docs',
|
||||
]),
|
||||
'title': 'State',
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite_aio].2
|
||||
dict({
|
||||
'$defs': dict({
|
||||
'InnerObject': dict({
|
||||
'properties': dict({
|
||||
'yo': dict({
|
||||
'title': 'Yo',
|
||||
'type': 'integer',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'yo',
|
||||
]),
|
||||
'title': 'InnerObject',
|
||||
'type': 'object',
|
||||
}),
|
||||
}),
|
||||
'properties': dict({
|
||||
'answer': dict({
|
||||
'anyOf': list([
|
||||
dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
dict({
|
||||
'type': 'null',
|
||||
}),
|
||||
]),
|
||||
'default': None,
|
||||
'title': 'Answer',
|
||||
}),
|
||||
'docs': dict({
|
||||
'items': dict({
|
||||
'type': 'string',
|
||||
}),
|
||||
'title': 'Docs',
|
||||
'type': 'array',
|
||||
}),
|
||||
'inner': dict({
|
||||
'$ref': '#/$defs/InnerObject',
|
||||
}),
|
||||
'query': dict({
|
||||
'title': 'Query',
|
||||
'type': 'string',
|
||||
}),
|
||||
}),
|
||||
'required': list([
|
||||
'query',
|
||||
'inner',
|
||||
'docs',
|
||||
]),
|
||||
'title': 'State',
|
||||
'type': 'object',
|
||||
})
|
||||
# ---
|
||||
# name: test_send_react_interrupt_control[memory]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
---
|
||||
config:
|
||||
flowchart:
|
||||
curve: linear
|
||||
---
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
agent(agent)
|
||||
foo([foo]):::last
|
||||
__start__ --> agent;
|
||||
agent -.-> foo;
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_send_react_interrupt_control[postgres_aio]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
agent(agent)
|
||||
foo([foo]):::last
|
||||
__start__ --> agent;
|
||||
agent -.-> foo;
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_send_react_interrupt_control[postgres_aio_pipe]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
agent(agent)
|
||||
foo([foo]):::last
|
||||
__start__ --> agent;
|
||||
agent -.-> foo;
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_send_react_interrupt_control[postgres_aio_pool]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
agent(agent)
|
||||
foo([foo]):::last
|
||||
__start__ --> agent;
|
||||
agent -.-> foo;
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_send_react_interrupt_control[postgres_aio_shallow]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
agent(agent)
|
||||
foo([foo]):::last
|
||||
__start__ --> agent;
|
||||
agent -.-> foo;
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_send_react_interrupt_control[sqlite_aio]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
agent(agent)
|
||||
foo([foo]):::last
|
||||
foo(foo)
|
||||
__end__([<p>__end__</p>]):::last
|
||||
__start__ --> agent;
|
||||
agent -.-> foo;
|
||||
foo --> __end__;
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
@@ -587,12 +587,10 @@ def test_conditional_graph(
|
||||
|
||||
app = workflow.compile()
|
||||
|
||||
if SHOULD_CHECK_SNAPSHOTS:
|
||||
if SHOULD_CHECK_SNAPSHOTS and checkpointer_name == "memory":
|
||||
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
|
||||
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
assert app.get_graph().draw_mermaid() == snapshot
|
||||
assert json.dumps(app.get_graph(xray=True).to_json(), indent=2) == snapshot
|
||||
assert app.get_graph(xray=True).draw_mermaid(with_styles=False) == snapshot
|
||||
|
||||
assert app.invoke({"input": "what is weather in sf"}) == {
|
||||
"input": "what is weather in sf",
|
||||
@@ -722,10 +720,6 @@ def test_conditional_graph(
|
||||
)
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
if SHOULD_CHECK_SNAPSHOTS:
|
||||
assert app_w_interrupt.get_graph().to_json() == snapshot
|
||||
assert app_w_interrupt.get_graph().draw_mermaid() == snapshot
|
||||
|
||||
assert [
|
||||
c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config)
|
||||
] == [
|
||||
@@ -1538,7 +1532,7 @@ def test_conditional_state_graph(
|
||||
|
||||
app = workflow.compile()
|
||||
|
||||
if SHOULD_CHECK_SNAPSHOTS:
|
||||
if SHOULD_CHECK_SNAPSHOTS and checkpointer_name == "memory":
|
||||
assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot
|
||||
assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot
|
||||
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
|
||||
@@ -3774,7 +3768,7 @@ def test_message_graph(
|
||||
# meaning you can use it as you would any other runnable
|
||||
app = workflow.compile()
|
||||
|
||||
if SHOULD_CHECK_SNAPSHOTS:
|
||||
if SHOULD_CHECK_SNAPSHOTS and checkpointer_name == "memory":
|
||||
assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot
|
||||
assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot
|
||||
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
|
||||
@@ -6234,10 +6228,13 @@ def test_start_branch_then(
|
||||
tool_two_graph.add_node("tool_two_slow", tool_two_slow)
|
||||
tool_two_graph.add_node("tool_two_fast", tool_two_fast)
|
||||
tool_two_graph.set_conditional_entry_point(
|
||||
lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast", then=END
|
||||
lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast",
|
||||
then=END,
|
||||
path_map=["tool_two_slow", "tool_two_fast"],
|
||||
)
|
||||
tool_two = tool_two_graph.compile()
|
||||
assert tool_two.get_graph().draw_mermaid() == snapshot
|
||||
if checkpointer_name == "memory":
|
||||
assert tool_two.get_graph().draw_mermaid() == snapshot
|
||||
|
||||
assert tool_two.invoke({"my_key": "value", "market": "DE"}) == {
|
||||
"my_key": "value slow",
|
||||
@@ -6516,6 +6513,7 @@ def test_branch_then(
|
||||
tool_two_graph.add_conditional_edges(
|
||||
source="prepare",
|
||||
path=lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast",
|
||||
path_map=["tool_two_slow", "tool_two_fast"],
|
||||
then="finish",
|
||||
)
|
||||
tool_two_graph.add_node("prepare", lambda s: {"my_key": " prepared"})
|
||||
@@ -6523,8 +6521,10 @@ def test_branch_then(
|
||||
tool_two_graph.add_node("tool_two_fast", lambda s: {"my_key": " fast"})
|
||||
tool_two_graph.add_node("finish", lambda s: {"my_key": " finished"})
|
||||
tool_two = tool_two_graph.compile()
|
||||
assert tool_two.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
assert tool_two.get_graph().draw_mermaid() == snapshot
|
||||
|
||||
if checkpointer_name == "memory":
|
||||
assert tool_two.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
assert tool_two.get_graph().draw_mermaid() == snapshot
|
||||
|
||||
assert tool_two.invoke({"my_key": "value", "market": "DE"}, debug=1) == {
|
||||
"my_key": "value prepared slow finished",
|
||||
@@ -9856,7 +9856,9 @@ def test_send_react_interrupt_control(
|
||||
builder.add_node(foo)
|
||||
builder.add_edge(START, "agent")
|
||||
graph = builder.compile()
|
||||
assert graph.get_graph().draw_mermaid() == snapshot
|
||||
|
||||
if checkpointer_name == "memory":
|
||||
assert graph.get_graph().draw_mermaid() == snapshot
|
||||
|
||||
assert graph.invoke({"messages": [HumanMessage("hello")]}) == {
|
||||
"messages": [
|
||||
@@ -10187,12 +10189,17 @@ def test_weather_subgraph(
|
||||
graph.add_node(normal_llm_node)
|
||||
graph.add_node("weather_graph", weather_graph)
|
||||
graph.add_edge(START, "router_node")
|
||||
graph.add_conditional_edges("router_node", route_after_prediction)
|
||||
graph.add_conditional_edges(
|
||||
"router_node",
|
||||
route_after_prediction,
|
||||
path_map=["weather_graph", "normal_llm_node"],
|
||||
)
|
||||
graph.add_edge("normal_llm_node", END)
|
||||
graph.add_edge("weather_graph", END)
|
||||
graph = graph.compile(checkpointer=checkpointer)
|
||||
|
||||
assert graph.get_graph(xray=1).draw_mermaid() == snapshot
|
||||
if checkpointer_name == "memory":
|
||||
assert graph.get_graph(xray=1).draw_mermaid() == snapshot
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
thread2 = {"configurable": {"thread_id": "2"}}
|
||||
|
||||
@@ -7041,7 +7041,11 @@ async def test_weather_subgraph(
|
||||
graph.add_node(normal_llm_node)
|
||||
graph.add_node("weather_graph", weather_graph)
|
||||
graph.add_edge(START, "router_node")
|
||||
graph.add_conditional_edges("router_node", route_after_prediction)
|
||||
graph.add_conditional_edges(
|
||||
"router_node",
|
||||
route_after_prediction,
|
||||
path_map=["weather_graph", "normal_llm_node"],
|
||||
)
|
||||
graph.add_edge("normal_llm_node", END)
|
||||
graph.add_edge("weather_graph", END)
|
||||
|
||||
@@ -7051,8 +7055,6 @@ async def test_weather_subgraph(
|
||||
async with awith_checkpointer(checkpointer_name) as checkpointer:
|
||||
graph = graph.compile(checkpointer=checkpointer)
|
||||
|
||||
assert graph.get_graph(xray=1).draw_mermaid() == snapshot
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
thread2 = {"configurable": {"thread_id": "2"}}
|
||||
inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]}
|
||||
|
||||
@@ -1,14 +1,9 @@
|
||||
import datetime
|
||||
import decimal
|
||||
import enum
|
||||
import functools
|
||||
import gc
|
||||
import ipaddress
|
||||
import json
|
||||
import logging
|
||||
import operator
|
||||
import pathlib
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
@@ -17,7 +12,6 @@ from collections import Counter, deque
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from random import randrange
|
||||
from typing import (
|
||||
Annotated,
|
||||
@@ -2153,7 +2147,7 @@ def test_conditional_entrypoint_to_multiple_state_graph(
|
||||
|
||||
workflow.add_node("get_weather", get_weather)
|
||||
workflow.add_edge("get_weather", END)
|
||||
workflow.set_conditional_entry_point(continue_to_weather)
|
||||
workflow.set_conditional_entry_point(continue_to_weather, path_map=["get_weather"])
|
||||
|
||||
app = workflow.compile()
|
||||
|
||||
@@ -2420,7 +2414,8 @@ def test_in_one_fan_out_state_graph_waiting_edge(
|
||||
|
||||
app = workflow.compile()
|
||||
|
||||
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
if checkpointer_name == "memory":
|
||||
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
|
||||
assert app.invoke({"query": "what is weather in sf"}) == {
|
||||
"query": "analyzed: query: what is weather in sf",
|
||||
@@ -2566,7 +2561,8 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch(
|
||||
|
||||
app = workflow.compile()
|
||||
|
||||
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
if checkpointer_name == "memory":
|
||||
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
|
||||
assert app.invoke({"query": "what is weather in sf"}, debug=True) == {
|
||||
"query": "analyzed: query: what is weather in sf",
|
||||
@@ -2716,9 +2712,10 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1(
|
||||
|
||||
app = workflow.compile()
|
||||
|
||||
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
assert app.get_input_jsonschema() == snapshot
|
||||
assert app.get_output_jsonschema() == snapshot
|
||||
if checkpointer_name == "memory":
|
||||
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
assert app.get_input_jsonschema() == snapshot
|
||||
assert app.get_output_jsonschema() == snapshot
|
||||
|
||||
with pytest.raises(ValidationError), assert_ctx_once():
|
||||
app.invoke({"query": {}})
|
||||
@@ -2906,7 +2903,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(
|
||||
|
||||
app = workflow.compile()
|
||||
|
||||
if SHOULD_CHECK_SNAPSHOTS:
|
||||
if SHOULD_CHECK_SNAPSHOTS and checkpointer_name == "memory":
|
||||
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
assert app.get_input_schema().model_json_schema() == snapshot
|
||||
assert app.get_output_schema().model_json_schema() == snapshot
|
||||
@@ -2970,8 +2967,6 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2(
|
||||
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
|
||||
def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input(
|
||||
snapshot: SnapshotAssertion,
|
||||
mocker: MockerFixture,
|
||||
request: pytest.FixtureRequest,
|
||||
checkpointer_name: str,
|
||||
) -> None:
|
||||
@@ -3101,328 +3096,6 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_inp
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("version", ["v1", "v2"])
|
||||
def test_nested_pydantic_models(version: str) -> None:
|
||||
"""Test that nested Pydantic models are properly constructed from leaf nodes up."""
|
||||
|
||||
# Define nested Pydantic models
|
||||
# Import necessary modules
|
||||
|
||||
if version == "v1":
|
||||
from pydantic.v1 import ( # type: ignore
|
||||
BaseModel,
|
||||
ByteSize,
|
||||
Field,
|
||||
SecretStr,
|
||||
confloat,
|
||||
conint,
|
||||
conlist,
|
||||
constr,
|
||||
)
|
||||
else:
|
||||
from pydantic import ( # type: ignore
|
||||
BaseModel,
|
||||
ByteSize,
|
||||
Field,
|
||||
SecretStr,
|
||||
confloat,
|
||||
conint,
|
||||
conlist,
|
||||
constr,
|
||||
)
|
||||
from pydantic.v1 import BaseModel as BaseModelV1
|
||||
|
||||
if BaseModel is BaseModelV1:
|
||||
pytest.skip("Cannot test pydantic v2 using installed version < 2")
|
||||
|
||||
class NestedModel(BaseModel):
|
||||
value: int
|
||||
name: str
|
||||
|
||||
# For constrained types
|
||||
PositiveInt = Annotated[int, Field(gt=0)]
|
||||
NonNegativeFloat = Annotated[float, Field(ge=0)]
|
||||
|
||||
# Enum type
|
||||
class UserRole(Enum):
|
||||
ADMIN = "admin"
|
||||
USER = "user"
|
||||
GUEST = "guest"
|
||||
|
||||
# Forward reference model
|
||||
class RecursiveModel(BaseModel):
|
||||
value: str
|
||||
child: Optional["RecursiveModel"] = None
|
||||
|
||||
# Discriminated union models
|
||||
class Cat(BaseModel):
|
||||
pet_type: Literal["cat"]
|
||||
meow: str
|
||||
|
||||
class Dog(BaseModel):
|
||||
pet_type: Literal["dog"]
|
||||
bark: str
|
||||
|
||||
# Cyclic reference model
|
||||
class Person(BaseModel):
|
||||
id: str
|
||||
name: str
|
||||
friends: list[str] = Field(default_factory=list) # IDs of friends
|
||||
|
||||
if version == "v2":
|
||||
conlist_type = conlist(item_type=int, min_length=2, max_length=5)
|
||||
else:
|
||||
conlist_type = conlist(item_type=int, min_items=2, max_items=5)
|
||||
|
||||
class State(BaseModel):
|
||||
# Basic nested model tests
|
||||
top_level: str
|
||||
auuid: uuid.UUID
|
||||
nested: NestedModel
|
||||
optional_nested: Annotated[Optional[NestedModel], lambda x, y: y, "Foo"]
|
||||
dict_nested: dict[str, NestedModel]
|
||||
simple_str_list: list[str]
|
||||
list_nested: Annotated[
|
||||
Union[dict, list[dict[str, NestedModel]]], lambda x, y: (x or []) + [y]
|
||||
]
|
||||
tuple_nested: tuple[str, NestedModel]
|
||||
tuple_list_nested: list[tuple[int, NestedModel]]
|
||||
complex_tuple: tuple[str, dict[str, tuple[int, NestedModel]]]
|
||||
|
||||
# Forward reference test
|
||||
recursive: RecursiveModel
|
||||
|
||||
# Discriminated union test
|
||||
pet: Union[Cat, Dog]
|
||||
|
||||
# Cyclic reference test
|
||||
people: dict[str, Person] # Map of ID -> Person
|
||||
|
||||
# Rich type adapters
|
||||
ip_address: ipaddress.IPv4Address
|
||||
ip_address_v6: ipaddress.IPv6Address
|
||||
amount: decimal.Decimal
|
||||
file_path: pathlib.Path
|
||||
timestamp: datetime.datetime
|
||||
date_only: datetime.date
|
||||
time_only: datetime.time
|
||||
duration: datetime.timedelta
|
||||
immutable_set: frozenset[int]
|
||||
binary_data: bytes
|
||||
pattern: re.Pattern
|
||||
secret: SecretStr
|
||||
file_size: ByteSize
|
||||
|
||||
# Constrained types
|
||||
positive_value: PositiveInt
|
||||
non_negative: NonNegativeFloat
|
||||
limited_string: constr(min_length=3, max_length=10)
|
||||
bounded_int: conint(ge=10, le=100)
|
||||
restricted_float: confloat(gt=0, lt=1)
|
||||
required_list: conlist_type
|
||||
|
||||
# Enum & Literal
|
||||
role: UserRole
|
||||
status: Literal["active", "inactive", "pending"]
|
||||
|
||||
# Annotated & NewType
|
||||
validated_age: Annotated[int, Field(gt=0, lt=120)]
|
||||
|
||||
# Generic containers with validators
|
||||
decimal_list: List[decimal.Decimal]
|
||||
id_tuple: tuple[uuid.UUID, uuid.UUID]
|
||||
|
||||
inputs = {
|
||||
# Basic nested models
|
||||
"top_level": "initial",
|
||||
"auuid": str(uuid.uuid4()),
|
||||
"nested": {"value": 42, "name": "test"},
|
||||
"optional_nested": {"value": 10, "name": "optional"},
|
||||
"dict_nested": {"a": {"value": 5, "name": "a"}},
|
||||
"list_nested": [{"a": {"value": 6, "name": "b"}}],
|
||||
"tuple_nested": ["tuple-key", {"value": 7, "name": "tuple-value"}],
|
||||
"tuple_list_nested": [[1, {"value": 8, "name": "tuple-in-list"}]],
|
||||
"simple_str_list": ["siss", "boom", "bah"],
|
||||
"complex_tuple": [
|
||||
"complex",
|
||||
{"nested": [9, {"value": 10, "name": "deep"}]},
|
||||
],
|
||||
# Forward reference
|
||||
"recursive": {"value": "parent", "child": {"value": "child", "child": None}},
|
||||
# Discriminated union (using a cat in this case)
|
||||
"pet": {"pet_type": "cat", "meow": "meow!"},
|
||||
# Cyclic references
|
||||
"people": {
|
||||
"1": {
|
||||
"id": "1",
|
||||
"name": "Alice",
|
||||
"friends": ["2", "3"], # Alice is friends with Bob and Charlie
|
||||
},
|
||||
"2": {
|
||||
"id": "2",
|
||||
"name": "Bob",
|
||||
"friends": ["1"], # Bob is friends with Alice
|
||||
},
|
||||
"3": {
|
||||
"id": "3",
|
||||
"name": "Charlie",
|
||||
"friends": ["1", "2"], # Charlie is friends with Alice and Bob
|
||||
},
|
||||
},
|
||||
# Rich type adapters
|
||||
"ip_address": "192.168.1.1",
|
||||
"ip_address_v6": "2001:db8::1",
|
||||
"amount": "123.45",
|
||||
"file_path": "/tmp/test.txt",
|
||||
"timestamp": "2025-04-07T10:58:04",
|
||||
"date_only": "2025-04-07",
|
||||
"time_only": "10:58:04",
|
||||
"duration": 3600, # seconds
|
||||
"immutable_set": [1, 2, 3, 4],
|
||||
"binary_data": b"hello world",
|
||||
"pattern": "^test$",
|
||||
"secret": "password123",
|
||||
"file_size": 1024,
|
||||
# Constrained types
|
||||
"positive_value": 42,
|
||||
"non_negative": 0.0,
|
||||
"limited_string": "test",
|
||||
"bounded_int": 50,
|
||||
"restricted_float": 0.5,
|
||||
"required_list": [10, 20, 30],
|
||||
# Enum & Literal
|
||||
"role": "admin",
|
||||
"status": "active",
|
||||
# Annotated & NewType
|
||||
"validated_age": 30,
|
||||
# Generic containers with validators
|
||||
"decimal_list": ["10.5", "20.75", "30.25"],
|
||||
"id_tuple": [str(uuid.uuid4()), str(uuid.uuid4())],
|
||||
}
|
||||
|
||||
update = {"top_level": "updated", "nested": {"value": 100, "name": "updated"}}
|
||||
|
||||
expected = State(**inputs)
|
||||
|
||||
def node_fn(state: State) -> dict:
|
||||
# Basic assertions
|
||||
assert isinstance(state.auuid, uuid.UUID)
|
||||
assert state == expected
|
||||
|
||||
# Rich type assertions
|
||||
assert isinstance(state.ip_address, ipaddress.IPv4Address)
|
||||
assert isinstance(state.ip_address_v6, ipaddress.IPv6Address)
|
||||
assert isinstance(state.amount, decimal.Decimal)
|
||||
assert isinstance(state.file_path, pathlib.Path)
|
||||
assert isinstance(state.timestamp, datetime.datetime)
|
||||
assert isinstance(state.date_only, datetime.date)
|
||||
assert isinstance(state.time_only, datetime.time)
|
||||
assert isinstance(state.duration, datetime.timedelta)
|
||||
assert isinstance(state.immutable_set, frozenset)
|
||||
assert isinstance(state.binary_data, bytes)
|
||||
assert isinstance(state.pattern, re.Pattern)
|
||||
|
||||
# Constrained types
|
||||
assert state.positive_value > 0
|
||||
assert state.non_negative >= 0
|
||||
assert 3 <= len(state.limited_string) <= 10
|
||||
assert 10 <= state.bounded_int <= 100
|
||||
assert 0 < state.restricted_float < 1
|
||||
assert 2 <= len(state.required_list) <= 5
|
||||
|
||||
# Enum & Literal
|
||||
assert state.role == UserRole.ADMIN
|
||||
assert state.status == "active"
|
||||
|
||||
# Annotated
|
||||
assert 0 < state.validated_age < 120
|
||||
|
||||
# Generic containers
|
||||
assert len(state.decimal_list) == 3
|
||||
assert len(state.id_tuple) == 2
|
||||
|
||||
return update
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("process", node_fn)
|
||||
builder.set_entry_point("process")
|
||||
builder.set_finish_point("process")
|
||||
graph = builder.compile()
|
||||
|
||||
result = graph.invoke(inputs.copy())
|
||||
|
||||
assert result == {**inputs, **update}
|
||||
|
||||
new_inputs = inputs.copy()
|
||||
new_inputs["list_nested"] = {"foo": "bar"}
|
||||
expected = State(**new_inputs)
|
||||
assert {**new_inputs, **update} == graph.invoke(new_inputs.copy())
|
||||
|
||||
|
||||
def test_pydantic_state_field_validator():
|
||||
from pydantic import BaseModel, field_validator, model_validator
|
||||
|
||||
class State(BaseModel):
|
||||
name: str
|
||||
text: str = ""
|
||||
only_root: int = 13
|
||||
|
||||
@field_validator("name", mode="after")
|
||||
@classmethod
|
||||
def validate_name(cls, value):
|
||||
if value[0].islower():
|
||||
raise ValueError("Name must start with a capital letter")
|
||||
return "Validated " + value
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def validate_amodel(cls, values: "State"):
|
||||
return values | {"only_root": 392}
|
||||
|
||||
input_state = {"name": "John"}
|
||||
|
||||
def process_node(state: State):
|
||||
assert State.model_validate(input_state) == state
|
||||
return {"text": "Hello, " + state.name + "!"}
|
||||
|
||||
builder = StateGraph(state_schema=State)
|
||||
builder.add_node("process", process_node)
|
||||
builder.add_edge(START, "process")
|
||||
builder.add_edge("process", END)
|
||||
g = builder.compile()
|
||||
res = g.invoke(input_state)
|
||||
assert res["text"] == "Hello, Validated John!"
|
||||
|
||||
|
||||
def test_pydantic_v1_state_root_validator():
|
||||
from pydantic.v1 import BaseModel, root_validator
|
||||
|
||||
class State(BaseModel):
|
||||
name: str
|
||||
text: str = ""
|
||||
only_root: int = 13
|
||||
|
||||
@root_validator(pre=True)
|
||||
@classmethod
|
||||
def validate(cls, values: dict):
|
||||
values["name"] = "Validated " + values["name"]
|
||||
return values | {"only_root": 396}
|
||||
|
||||
input_state = {"name": "John"}
|
||||
|
||||
def process_node(state: State):
|
||||
assert State(**input_state) == state
|
||||
return {"text": "Hello, " + state.name + "!"}
|
||||
|
||||
builder = StateGraph(state_schema=State)
|
||||
builder.add_node("process", process_node)
|
||||
builder.add_edge(START, "process")
|
||||
builder.add_edge("process", END)
|
||||
g = builder.compile()
|
||||
res = g.invoke(input_state)
|
||||
assert res["text"] == "Hello, Validated John!"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
|
||||
def test_in_one_fan_out_state_graph_waiting_edge_plus_regular(
|
||||
request: pytest.FixtureRequest, checkpointer_name: str
|
||||
@@ -4804,7 +4477,9 @@ def test_xray_lance(snapshot: SnapshotAssertion):
|
||||
# Flow
|
||||
interview_builder.add_edge(START, "ask_question")
|
||||
interview_builder.add_edge("ask_question", "answer_question")
|
||||
interview_builder.add_conditional_edges("answer_question", route_messages)
|
||||
interview_builder.add_conditional_edges(
|
||||
"answer_question", route_messages, ["ask_question", END]
|
||||
)
|
||||
|
||||
# Set up memory
|
||||
memory = InMemorySaver()
|
||||
@@ -7336,6 +7011,8 @@ def test_node_destinations() -> None:
|
||||
Edge(source="__start__", target="child", data=None, conditional=False),
|
||||
Edge(source="child", target="node_b", data=None, conditional=True),
|
||||
Edge(source="child", target="node_c", data=None, conditional=True),
|
||||
Edge(source="node_b", target="__end__", data=None, conditional=False),
|
||||
Edge(source="node_c", target="__end__", data=None, conditional=False),
|
||||
] == graph.edges
|
||||
|
||||
# destinations w/ dicts
|
||||
@@ -7354,6 +7031,8 @@ def test_node_destinations() -> None:
|
||||
Edge(source="__start__", target="child", data=None, conditional=False),
|
||||
Edge(source="child", target="node_b", data="foo", conditional=True),
|
||||
Edge(source="child", target="node_c", data="bar", conditional=True),
|
||||
Edge(source="node_b", target="__end__", data=None, conditional=False),
|
||||
Edge(source="node_c", target="__end__", data=None, conditional=False),
|
||||
] == graph.edges
|
||||
|
||||
|
||||
|
||||
@@ -3610,7 +3610,8 @@ async def test_send_react_interrupt_control(
|
||||
builder.add_node(foo)
|
||||
builder.add_edge(START, "agent")
|
||||
graph = builder.compile()
|
||||
assert graph.get_graph().draw_mermaid() == snapshot
|
||||
if checkpointer_name == "memory":
|
||||
assert graph.get_graph().draw_mermaid() == snapshot
|
||||
|
||||
assert await graph.ainvoke({"messages": [HumanMessage("hello")]}) == {
|
||||
"messages": [
|
||||
@@ -3928,22 +3929,29 @@ async def test_max_concurrency_control(checkpointer_name: str) -> None:
|
||||
builder.add_edge(START, "1")
|
||||
graph = builder.compile()
|
||||
|
||||
assert (
|
||||
graph.get_graph().draw_mermaid()
|
||||
== """%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
if checkpointer_name == "memory":
|
||||
assert (
|
||||
graph.get_graph().draw_mermaid()
|
||||
== """---
|
||||
config:
|
||||
flowchart:
|
||||
curve: linear
|
||||
---
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
1(1)
|
||||
2(2)
|
||||
3([3]):::last
|
||||
__start__ --> 1;
|
||||
3(3)
|
||||
__end__([<p>__end__</p>]):::last
|
||||
1 -.-> 2;
|
||||
2 -.-> 3;
|
||||
__start__ --> 1;
|
||||
3 --> __end__;
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
"""
|
||||
)
|
||||
)
|
||||
|
||||
assert await graph.ainvoke(["0"], debug=True) == ["0", "1", *range(100), "3"]
|
||||
assert node2_max_currently == 100
|
||||
@@ -4980,7 +4988,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydant
|
||||
|
||||
app = workflow.compile()
|
||||
|
||||
if SHOULD_CHECK_SNAPSHOTS:
|
||||
if SHOULD_CHECK_SNAPSHOTS and checkpointer_name == "memory":
|
||||
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
assert app.get_input_schema().model_json_schema() == snapshot
|
||||
assert app.get_output_schema().model_json_schema() == snapshot
|
||||
|
||||
@@ -1,14 +1,26 @@
|
||||
import datetime
|
||||
import decimal
|
||||
import ipaddress
|
||||
import pathlib
|
||||
import re
|
||||
import sys
|
||||
import typing
|
||||
import uuid
|
||||
from enum import Enum
|
||||
from typing import Annotated, List, Literal, Optional, Union
|
||||
|
||||
import pydantic
|
||||
import typing_extensions
|
||||
import pytest
|
||||
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.graph.state import StateGraph
|
||||
from langgraph.utils.pydantic import is_supported_by_pydantic
|
||||
|
||||
|
||||
def test_is_supported_by_pydantic() -> None:
|
||||
"""Test if types are supported by pydantic."""
|
||||
import typing
|
||||
|
||||
import pydantic
|
||||
import typing_extensions
|
||||
|
||||
class TypedDictExtensions(typing_extensions.TypedDict):
|
||||
x: int
|
||||
@@ -41,3 +53,325 @@ def test_is_supported_by_pydantic() -> None:
|
||||
assert is_supported_by_pydantic(PydanticModelV1) is False
|
||||
|
||||
assert is_supported_by_pydantic(int) is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize("version", ["v1", "v2"])
|
||||
def test_nested_pydantic_models(version: str) -> None:
|
||||
"""Test that nested Pydantic models are properly constructed from leaf nodes up."""
|
||||
|
||||
# Define nested Pydantic models
|
||||
# Import necessary modules
|
||||
|
||||
if version == "v1":
|
||||
from pydantic.v1 import ( # type: ignore
|
||||
BaseModel,
|
||||
ByteSize,
|
||||
Field,
|
||||
SecretStr,
|
||||
confloat,
|
||||
conint,
|
||||
conlist,
|
||||
constr,
|
||||
)
|
||||
else:
|
||||
from pydantic import ( # type: ignore
|
||||
BaseModel,
|
||||
ByteSize,
|
||||
Field,
|
||||
SecretStr,
|
||||
confloat,
|
||||
conint,
|
||||
conlist,
|
||||
constr,
|
||||
)
|
||||
from pydantic.v1 import BaseModel as BaseModelV1
|
||||
|
||||
if BaseModel is BaseModelV1:
|
||||
pytest.skip("Cannot test pydantic v2 using installed version < 2")
|
||||
|
||||
class NestedModel(BaseModel):
|
||||
value: int
|
||||
name: str
|
||||
|
||||
# For constrained types
|
||||
PositiveInt = Annotated[int, Field(gt=0)]
|
||||
NonNegativeFloat = Annotated[float, Field(ge=0)]
|
||||
|
||||
# Enum type
|
||||
class UserRole(Enum):
|
||||
ADMIN = "admin"
|
||||
USER = "user"
|
||||
GUEST = "guest"
|
||||
|
||||
# Forward reference model
|
||||
class RecursiveModel(BaseModel):
|
||||
value: str
|
||||
child: Optional["RecursiveModel"] = None
|
||||
|
||||
# Discriminated union models
|
||||
class Cat(BaseModel):
|
||||
pet_type: Literal["cat"]
|
||||
meow: str
|
||||
|
||||
class Dog(BaseModel):
|
||||
pet_type: Literal["dog"]
|
||||
bark: str
|
||||
|
||||
# Cyclic reference model
|
||||
class Person(BaseModel):
|
||||
id: str
|
||||
name: str
|
||||
friends: list[str] = Field(default_factory=list) # IDs of friends
|
||||
|
||||
if version == "v2":
|
||||
conlist_type = conlist(item_type=int, min_length=2, max_length=5)
|
||||
else:
|
||||
conlist_type = conlist(item_type=int, min_items=2, max_items=5)
|
||||
|
||||
class State(BaseModel):
|
||||
# Basic nested model tests
|
||||
top_level: str
|
||||
auuid: uuid.UUID
|
||||
nested: NestedModel
|
||||
optional_nested: Annotated[Optional[NestedModel], lambda x, y: y, "Foo"]
|
||||
dict_nested: dict[str, NestedModel]
|
||||
simple_str_list: list[str]
|
||||
list_nested: Annotated[
|
||||
Union[dict, list[dict[str, NestedModel]]], lambda x, y: (x or []) + [y]
|
||||
]
|
||||
tuple_nested: tuple[str, NestedModel]
|
||||
tuple_list_nested: list[tuple[int, NestedModel]]
|
||||
complex_tuple: tuple[str, dict[str, tuple[int, NestedModel]]]
|
||||
|
||||
# Forward reference test
|
||||
recursive: RecursiveModel
|
||||
|
||||
# Discriminated union test
|
||||
pet: Union[Cat, Dog]
|
||||
|
||||
# Cyclic reference test
|
||||
people: dict[str, Person] # Map of ID -> Person
|
||||
|
||||
# Rich type adapters
|
||||
ip_address: ipaddress.IPv4Address
|
||||
ip_address_v6: ipaddress.IPv6Address
|
||||
amount: decimal.Decimal
|
||||
file_path: pathlib.Path
|
||||
timestamp: datetime.datetime
|
||||
date_only: datetime.date
|
||||
time_only: datetime.time
|
||||
duration: datetime.timedelta
|
||||
immutable_set: frozenset[int]
|
||||
binary_data: bytes
|
||||
pattern: re.Pattern
|
||||
secret: SecretStr
|
||||
file_size: ByteSize
|
||||
|
||||
# Constrained types
|
||||
positive_value: PositiveInt
|
||||
non_negative: NonNegativeFloat
|
||||
limited_string: constr(min_length=3, max_length=10)
|
||||
bounded_int: conint(ge=10, le=100)
|
||||
restricted_float: confloat(gt=0, lt=1)
|
||||
required_list: conlist_type
|
||||
|
||||
# Enum & Literal
|
||||
role: UserRole
|
||||
status: Literal["active", "inactive", "pending"]
|
||||
|
||||
# Annotated & NewType
|
||||
validated_age: Annotated[int, Field(gt=0, lt=120)]
|
||||
|
||||
# Generic containers with validators
|
||||
decimal_list: List[decimal.Decimal]
|
||||
id_tuple: tuple[uuid.UUID, uuid.UUID]
|
||||
|
||||
inputs = {
|
||||
# Basic nested models
|
||||
"top_level": "initial",
|
||||
"auuid": str(uuid.uuid4()),
|
||||
"nested": {"value": 42, "name": "test"},
|
||||
"optional_nested": {"value": 10, "name": "optional"},
|
||||
"dict_nested": {"a": {"value": 5, "name": "a"}},
|
||||
"list_nested": [{"a": {"value": 6, "name": "b"}}],
|
||||
"tuple_nested": ["tuple-key", {"value": 7, "name": "tuple-value"}],
|
||||
"tuple_list_nested": [[1, {"value": 8, "name": "tuple-in-list"}]],
|
||||
"simple_str_list": ["siss", "boom", "bah"],
|
||||
"complex_tuple": [
|
||||
"complex",
|
||||
{"nested": [9, {"value": 10, "name": "deep"}]},
|
||||
],
|
||||
# Forward reference
|
||||
"recursive": {"value": "parent", "child": {"value": "child", "child": None}},
|
||||
# Discriminated union (using a cat in this case)
|
||||
"pet": {"pet_type": "cat", "meow": "meow!"},
|
||||
# Cyclic references
|
||||
"people": {
|
||||
"1": {
|
||||
"id": "1",
|
||||
"name": "Alice",
|
||||
"friends": ["2", "3"], # Alice is friends with Bob and Charlie
|
||||
},
|
||||
"2": {
|
||||
"id": "2",
|
||||
"name": "Bob",
|
||||
"friends": ["1"], # Bob is friends with Alice
|
||||
},
|
||||
"3": {
|
||||
"id": "3",
|
||||
"name": "Charlie",
|
||||
"friends": ["1", "2"], # Charlie is friends with Alice and Bob
|
||||
},
|
||||
},
|
||||
# Rich type adapters
|
||||
"ip_address": "192.168.1.1",
|
||||
"ip_address_v6": "2001:db8::1",
|
||||
"amount": "123.45",
|
||||
"file_path": "/tmp/test.txt",
|
||||
"timestamp": "2025-04-07T10:58:04",
|
||||
"date_only": "2025-04-07",
|
||||
"time_only": "10:58:04",
|
||||
"duration": 3600, # seconds
|
||||
"immutable_set": [1, 2, 3, 4],
|
||||
"binary_data": b"hello world",
|
||||
"pattern": "^test$",
|
||||
"secret": "password123",
|
||||
"file_size": 1024,
|
||||
# Constrained types
|
||||
"positive_value": 42,
|
||||
"non_negative": 0.0,
|
||||
"limited_string": "test",
|
||||
"bounded_int": 50,
|
||||
"restricted_float": 0.5,
|
||||
"required_list": [10, 20, 30],
|
||||
# Enum & Literal
|
||||
"role": "admin",
|
||||
"status": "active",
|
||||
# Annotated & NewType
|
||||
"validated_age": 30,
|
||||
# Generic containers with validators
|
||||
"decimal_list": ["10.5", "20.75", "30.25"],
|
||||
"id_tuple": [str(uuid.uuid4()), str(uuid.uuid4())],
|
||||
}
|
||||
|
||||
update = {"top_level": "updated", "nested": {"value": 100, "name": "updated"}}
|
||||
|
||||
expected = State(**inputs)
|
||||
|
||||
def node_fn(state: State) -> dict:
|
||||
# Basic assertions
|
||||
assert isinstance(state.auuid, uuid.UUID)
|
||||
assert state == expected
|
||||
|
||||
# Rich type assertions
|
||||
assert isinstance(state.ip_address, ipaddress.IPv4Address)
|
||||
assert isinstance(state.ip_address_v6, ipaddress.IPv6Address)
|
||||
assert isinstance(state.amount, decimal.Decimal)
|
||||
assert isinstance(state.file_path, pathlib.Path)
|
||||
assert isinstance(state.timestamp, datetime.datetime)
|
||||
assert isinstance(state.date_only, datetime.date)
|
||||
assert isinstance(state.time_only, datetime.time)
|
||||
assert isinstance(state.duration, datetime.timedelta)
|
||||
assert isinstance(state.immutable_set, frozenset)
|
||||
assert isinstance(state.binary_data, bytes)
|
||||
assert isinstance(state.pattern, re.Pattern)
|
||||
|
||||
# Constrained types
|
||||
assert state.positive_value > 0
|
||||
assert state.non_negative >= 0
|
||||
assert 3 <= len(state.limited_string) <= 10
|
||||
assert 10 <= state.bounded_int <= 100
|
||||
assert 0 < state.restricted_float < 1
|
||||
assert 2 <= len(state.required_list) <= 5
|
||||
|
||||
# Enum & Literal
|
||||
assert state.role == UserRole.ADMIN
|
||||
assert state.status == "active"
|
||||
|
||||
# Annotated
|
||||
assert 0 < state.validated_age < 120
|
||||
|
||||
# Generic containers
|
||||
assert len(state.decimal_list) == 3
|
||||
assert len(state.id_tuple) == 2
|
||||
|
||||
return update
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("process", node_fn)
|
||||
builder.set_entry_point("process")
|
||||
builder.set_finish_point("process")
|
||||
graph = builder.compile()
|
||||
|
||||
result = graph.invoke(inputs.copy())
|
||||
|
||||
assert result == {**inputs, **update}
|
||||
|
||||
new_inputs = inputs.copy()
|
||||
new_inputs["list_nested"] = {"foo": "bar"}
|
||||
expected = State(**new_inputs)
|
||||
assert {**new_inputs, **update} == graph.invoke(new_inputs.copy())
|
||||
|
||||
|
||||
def test_pydantic_state_field_validator():
|
||||
from pydantic import BaseModel, field_validator, model_validator
|
||||
|
||||
class State(BaseModel):
|
||||
name: str
|
||||
text: str = ""
|
||||
only_root: int = 13
|
||||
|
||||
@field_validator("name", mode="after")
|
||||
@classmethod
|
||||
def validate_name(cls, value):
|
||||
if value[0].islower():
|
||||
raise ValueError("Name must start with a capital letter")
|
||||
return "Validated " + value
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def validate_amodel(cls, values: "State"):
|
||||
return values | {"only_root": 392}
|
||||
|
||||
input_state = {"name": "John"}
|
||||
|
||||
def process_node(state: State):
|
||||
assert State.model_validate(input_state) == state
|
||||
return {"text": "Hello, " + state.name + "!"}
|
||||
|
||||
builder = StateGraph(state_schema=State)
|
||||
builder.add_node("process", process_node)
|
||||
builder.add_edge(START, "process")
|
||||
builder.add_edge("process", END)
|
||||
g = builder.compile()
|
||||
res = g.invoke(input_state)
|
||||
assert res["text"] == "Hello, Validated John!"
|
||||
|
||||
|
||||
def test_pydantic_v1_state_root_validator():
|
||||
from pydantic.v1 import BaseModel, root_validator
|
||||
|
||||
class State(BaseModel):
|
||||
name: str
|
||||
text: str = ""
|
||||
only_root: int = 13
|
||||
|
||||
@root_validator(pre=True)
|
||||
@classmethod
|
||||
def validate(cls, values: dict):
|
||||
values["name"] = "Validated " + values["name"]
|
||||
return values | {"only_root": 396}
|
||||
|
||||
input_state = {"name": "John"}
|
||||
|
||||
def process_node(state: State):
|
||||
assert State(**input_state) == state
|
||||
return {"text": "Hello, " + state.name + "!"}
|
||||
|
||||
builder = StateGraph(state_schema=State)
|
||||
builder.add_node("process", process_node)
|
||||
builder.add_edge(START, "process")
|
||||
builder.add_edge("process", END)
|
||||
g = builder.compile()
|
||||
res = g.invoke(input_state)
|
||||
assert res["text"] == "Hello, Validated John!"
|
||||
|
||||
Reference in New Issue
Block a user