mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-20 06:35:46 +02:00
Rewrite graph drawing logic
- 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 commit is contained in:
Nuno Campos
@@ -29,12 +29,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 +129,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 +141,8 @@ class Branch(NamedTuple):
|
||||
name=None,
|
||||
trace=False,
|
||||
func_accepts_config=True,
|
||||
)
|
||||
),
|
||||
writer(list(self.ends.values())) if self.ends else None,
|
||||
)
|
||||
|
||||
def _route(
|
||||
@@ -147,9 +151,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 +174,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 +193,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 +210,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,
|
||||
declared=_control_branch(Command(goto=tuple(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 = [
|
||||
(
|
||||
@@ -861,9 +866,8 @@ class CompiledStateGraph(CompiledGraph):
|
||||
),
|
||||
)
|
||||
)
|
||||
ChannelWrite.do_write(
|
||||
config, cast(Sequence[Union[Send, ChannelWriteEntry]], writes)
|
||||
)
|
||||
return writes
|
||||
return []
|
||||
|
||||
if with_reader:
|
||||
# get schema
|
||||
@@ -891,7 +895,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:
|
||||
|
||||
@@ -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,207 @@
|
||||
from collections import defaultdict
|
||||
from typing import Any, Mapping, Optional, Sequence, Union
|
||||
|
||||
from langchain_core.runnables.config import RunnableConfig
|
||||
from langchain_core.runnables.graph import Graph
|
||||
|
||||
from langgraph.channels.base import BaseChannel
|
||||
from langgraph.checkpoint.base import BaseCheckpointSaver
|
||||
from langgraph.constants import CONF, CONFIG_KEY_SEND, END, INPUT
|
||||
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 DEFAULT_BOUND, PregelNode
|
||||
from langgraph.pregel.write import ChannelWrite, ChannelWriteTupleEntry
|
||||
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)
|
||||
edges: list[tuple[str, str, bool]] = []
|
||||
|
||||
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):
|
||||
declared_seen: set[Any] = set()
|
||||
sources: dict[str, set[tuple[str, bool]]] = {}
|
||||
step_sources: dict[str, set[tuple[str, bool]]] = {}
|
||||
# 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,
|
||||
)
|
||||
# run the pregel loop
|
||||
while tasks:
|
||||
conditionals = set()
|
||||
# run task writers
|
||||
for task in tasks.values():
|
||||
for w in task.writers:
|
||||
if isinstance(w, ChannelWrite):
|
||||
w.invoke(None, task.config)
|
||||
# apply declared writes (Command)
|
||||
for entry in w.writes:
|
||||
if (
|
||||
isinstance(entry, ChannelWriteTupleEntry)
|
||||
and entry.declared
|
||||
and entry not in conditionals
|
||||
):
|
||||
# visit only once
|
||||
declared_seen.add(entry)
|
||||
# apply them
|
||||
current_len = len(task.writes)
|
||||
task.config[CONF][CONFIG_KEY_SEND](entry.declared)
|
||||
conditionals.update(list(task.writes)[current_len:])
|
||||
elif w not in declared_seen:
|
||||
# visit only once
|
||||
declared_seen.add(w)
|
||||
# get declared writes
|
||||
if writes := ChannelWrite.get_declared_writes(w):
|
||||
# apply them
|
||||
current_len = len(task.writes)
|
||||
ChannelWrite.do_write(task.config, writes)
|
||||
conditionals.update(list(task.writes)[current_len:])
|
||||
# collect sources
|
||||
step_sources = {
|
||||
task.name: {(w[0], w in conditionals) 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]]] = defaultdict(set)
|
||||
for src, triggers in sources.items():
|
||||
for trigger, cond in triggers:
|
||||
trigger_to_sources[trigger].add((src, cond))
|
||||
# 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 in sorted(trigger_to_sources[trigger]):
|
||||
edges.append((src, task.name, cond))
|
||||
# assemble the graph
|
||||
graph = Graph()
|
||||
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)
|
||||
for src, dest, is_conditional in edges:
|
||||
# TODO conditional labels
|
||||
graph.add_edge(
|
||||
graph.nodes[src], graph.nodes[dest], conditional=is_conditional
|
||||
)
|
||||
# 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=last)
|
||||
elif edge.target == name:
|
||||
graph.edges[idx] = edge.copy(target=first)
|
||||
# add end edges
|
||||
if step_sources:
|
||||
end = graph.add_node(DEFAULT_BOUND, END)
|
||||
for src in step_sources:
|
||||
graph.add_edge(graph.nodes[src], end)
|
||||
termini = set(d for _, d, _ in edges).difference((s for s, _, _ in edges))
|
||||
for src in termini.union(step_sources):
|
||||
# TODO conditional labels
|
||||
graph.add_edge(graph.nodes[src], end, conditional=src not in termini)
|
||||
|
||||
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."""
|
||||
declared: Optional[Sequence[tuple[str, Any]]] = 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,10 +133,10 @@ 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]] = []
|
||||
@@ -162,14 +165,26 @@ class ChannelWrite(RunnableCallable):
|
||||
"""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_declared_writes(
|
||||
runnable: Runnable,
|
||||
) -> Optional[Sequence[Union[ChannelWriteEntry, Send]]]:
|
||||
"""Used to get the writes a writer declares for static analysis."""
|
||||
if writes := getattr(runnable, "_is_channel_writer", MISSING):
|
||||
return writes if writes is not MISSING else None
|
||||
|
||||
@staticmethod
|
||||
def register_writer(
|
||||
runnable: R,
|
||||
declared: Optional[Sequence[Union[ChannelWriteEntry, Send]]] = 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", declared)
|
||||
return runnable
|
||||
|
||||
@@ -241,11 +241,6 @@
|
||||
'''
|
||||
{
|
||||
"nodes": [
|
||||
{
|
||||
"id": "__start__",
|
||||
"type": "schema",
|
||||
"data": "__start__"
|
||||
},
|
||||
{
|
||||
"id": "agent",
|
||||
"type": "runnable",
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -3,12 +3,12 @@
|
||||
'''
|
||||
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;
|
||||
analyzer_one --> retriever_one;
|
||||
retriever_one --> qa;
|
||||
retriever_two --> qa;
|
||||
qa --> __end__;
|
||||
|
||||
'''
|
||||
# ---
|
||||
@@ -120,611 +120,6 @@
|
||||
'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'}}}%%
|
||||
@@ -740,78 +135,3 @@
|
||||
|
||||
'''
|
||||
# ---
|
||||
# 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
|
||||
__start__ --> agent;
|
||||
agent -.-> foo;
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
|
||||
@@ -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,
|
||||
@@ -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
|
||||
|
||||
@@ -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,9 +3929,10 @@ 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()
|
||||
== """%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([<p>__start__</p>]):::first
|
||||
1(1)
|
||||
@@ -3943,7 +3945,7 @@ graph TD;
|
||||
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 +4982,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