diff --git a/libs/langgraph/langgraph/graph/branch.py b/libs/langgraph/langgraph/graph/branch.py
index 33a2aca1e..1c9bdad53 100644
--- a/libs/langgraph/langgraph/graph/branch.py
+++ b/libs/langgraph/langgraph/graph/branch.py
@@ -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
diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py
index fa28243fb..9f8b2b785 100644
--- a/libs/langgraph/langgraph/graph/graph.py
+++ b/libs/langgraph/langgraph/graph/graph.py
@@ -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(),
- }
diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py
index fed12fd8b..ff2b9e3ff 100644
--- a/libs/langgraph/langgraph/graph/state.py
+++ b/libs/langgraph/langgraph/graph/state.py
@@ -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:
diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py
index 2b016e587..231144794 100644
--- a/libs/langgraph/langgraph/pregel/__init__.py
+++ b/libs/langgraph/langgraph/pregel/__init__.py
@@ -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 {})}
diff --git a/libs/langgraph/langgraph/pregel/draw.py b/libs/langgraph/langgraph/pregel/draw.py
new file mode 100644
index 000000000..c2e652106
--- /dev/null
+++ b/libs/langgraph/langgraph/pregel/draw.py
@@ -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
diff --git a/libs/langgraph/langgraph/pregel/write.py b/libs/langgraph/langgraph/pregel/write.py
index 234c1f5d7..21a8654fe 100644
--- a/libs/langgraph/langgraph/pregel/write.py
+++ b/libs/langgraph/langgraph/pregel/write.py
@@ -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
diff --git a/libs/langgraph/tests/__snapshots__/test_large_cases.ambr b/libs/langgraph/tests/__snapshots__/test_large_cases.ambr
index adb41688d..64671e4db 100644
--- a/libs/langgraph/tests/__snapshots__/test_large_cases.ambr
+++ b/libs/langgraph/tests/__snapshots__/test_large_cases.ambr
@@ -241,11 +241,6 @@
'''
{
"nodes": [
- {
- "id": "__start__",
- "type": "schema",
- "data": "__start__"
- },
{
"id": "agent",
"type": "runnable",
diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr
index 336d20bf2..c37ac91dc 100644
--- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr
+++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr
@@ -303,90 +303,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;
-
- '''
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge[postgres]
- '''
- 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[postgres_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[postgres_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[postgres_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[sqlite]
- '''
- 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[sqlite_aes]
- '''
- 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;
'''
# ---
@@ -394,12 +316,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__;
'''
# ---
@@ -460,436 +382,16 @@
'type': 'object',
})
# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres]
- '''
- 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_pydantic1[postgres].1
- dict({
- 'definitions': dict({
- 'InnerObject': dict({
- 'properties': dict({
- 'yo': dict({
- 'title': 'Yo',
- 'type': 'integer',
- }),
- }),
- 'required': list([
- 'yo',
- ]),
- 'title': 'InnerObject',
- 'type': 'object',
- }),
- }),
- 'properties': dict({
- 'inner': dict({
- '$ref': '#/definitions/InnerObject',
- }),
- 'query': dict({
- 'title': 'Query',
- 'type': 'string',
- }),
- }),
- 'required': list([
- 'query',
- 'inner',
- ]),
- 'title': 'Input',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres].2
- dict({
- 'properties': dict({
- 'answer': dict({
- 'title': 'Answer',
- 'type': 'string',
- }),
- 'docs': dict({
- 'items': dict({
- 'type': 'string',
- }),
- 'title': 'Docs',
- 'type': 'array',
- }),
- }),
- 'required': list([
- 'answer',
- 'docs',
- ]),
- 'title': 'Output',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_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_pydantic1[postgres_pipe].1
- dict({
- 'definitions': dict({
- 'InnerObject': dict({
- 'properties': dict({
- 'yo': dict({
- 'title': 'Yo',
- 'type': 'integer',
- }),
- }),
- 'required': list([
- 'yo',
- ]),
- 'title': 'InnerObject',
- 'type': 'object',
- }),
- }),
- 'properties': dict({
- 'inner': dict({
- '$ref': '#/definitions/InnerObject',
- }),
- 'query': dict({
- 'title': 'Query',
- 'type': 'string',
- }),
- }),
- 'required': list([
- 'query',
- 'inner',
- ]),
- 'title': 'Input',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_pipe].2
- dict({
- 'properties': dict({
- 'answer': dict({
- 'title': 'Answer',
- 'type': 'string',
- }),
- 'docs': dict({
- 'items': dict({
- 'type': 'string',
- }),
- 'title': 'Docs',
- 'type': 'array',
- }),
- }),
- 'required': list([
- 'answer',
- 'docs',
- ]),
- 'title': 'Output',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_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_pydantic1[postgres_pool].1
- dict({
- 'definitions': dict({
- 'InnerObject': dict({
- 'properties': dict({
- 'yo': dict({
- 'title': 'Yo',
- 'type': 'integer',
- }),
- }),
- 'required': list([
- 'yo',
- ]),
- 'title': 'InnerObject',
- 'type': 'object',
- }),
- }),
- 'properties': dict({
- 'inner': dict({
- '$ref': '#/definitions/InnerObject',
- }),
- 'query': dict({
- 'title': 'Query',
- 'type': 'string',
- }),
- }),
- 'required': list([
- 'query',
- 'inner',
- ]),
- 'title': 'Input',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_pool].2
- dict({
- 'properties': dict({
- 'answer': dict({
- 'title': 'Answer',
- 'type': 'string',
- }),
- 'docs': dict({
- 'items': dict({
- 'type': 'string',
- }),
- 'title': 'Docs',
- 'type': 'array',
- }),
- }),
- 'required': list([
- 'answer',
- 'docs',
- ]),
- 'title': 'Output',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_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_pydantic1[postgres_shallow].1
- dict({
- 'definitions': dict({
- 'InnerObject': dict({
- 'properties': dict({
- 'yo': dict({
- 'title': 'Yo',
- 'type': 'integer',
- }),
- }),
- 'required': list([
- 'yo',
- ]),
- 'title': 'InnerObject',
- 'type': 'object',
- }),
- }),
- 'properties': dict({
- 'inner': dict({
- '$ref': '#/definitions/InnerObject',
- }),
- 'query': dict({
- 'title': 'Query',
- 'type': 'string',
- }),
- }),
- 'required': list([
- 'query',
- 'inner',
- ]),
- 'title': 'Input',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[postgres_shallow].2
- dict({
- 'properties': dict({
- 'answer': dict({
- 'title': 'Answer',
- 'type': 'string',
- }),
- 'docs': dict({
- 'items': dict({
- 'type': 'string',
- }),
- 'title': 'Docs',
- 'type': 'array',
- }),
- }),
- 'required': list([
- 'answer',
- 'docs',
- ]),
- 'title': 'Output',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[sqlite]
- '''
- 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_pydantic1[sqlite].1
- dict({
- 'definitions': dict({
- 'InnerObject': dict({
- 'properties': dict({
- 'yo': dict({
- 'title': 'Yo',
- 'type': 'integer',
- }),
- }),
- 'required': list([
- 'yo',
- ]),
- 'title': 'InnerObject',
- 'type': 'object',
- }),
- }),
- 'properties': dict({
- 'inner': dict({
- '$ref': '#/definitions/InnerObject',
- }),
- 'query': dict({
- 'title': 'Query',
- 'type': 'string',
- }),
- }),
- 'required': list([
- 'query',
- 'inner',
- ]),
- 'title': 'Input',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[sqlite].2
- dict({
- 'properties': dict({
- 'answer': dict({
- 'title': 'Answer',
- 'type': 'string',
- }),
- 'docs': dict({
- 'items': dict({
- 'type': 'string',
- }),
- 'title': 'Docs',
- 'type': 'array',
- }),
- }),
- 'required': list([
- 'answer',
- 'docs',
- ]),
- 'title': 'Output',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[sqlite_aes]
- '''
- 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_pydantic1[sqlite_aes].1
- dict({
- 'definitions': dict({
- 'InnerObject': dict({
- 'properties': dict({
- 'yo': dict({
- 'title': 'Yo',
- 'type': 'integer',
- }),
- }),
- 'required': list([
- 'yo',
- ]),
- 'title': 'InnerObject',
- 'type': 'object',
- }),
- }),
- 'properties': dict({
- 'inner': dict({
- '$ref': '#/definitions/InnerObject',
- }),
- 'query': dict({
- 'title': 'Query',
- 'type': 'string',
- }),
- }),
- 'required': list([
- 'query',
- 'inner',
- ]),
- 'title': 'Input',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1[sqlite_aes].2
- dict({
- 'properties': dict({
- 'answer': dict({
- 'title': 'Answer',
- 'type': 'string',
- }),
- 'docs': dict({
- 'items': dict({
- 'type': 'string',
- }),
- 'title': 'Docs',
- 'type': 'array',
- }),
- }),
- 'required': list([
- 'answer',
- 'docs',
- ]),
- 'title': 'Output',
- 'type': 'object',
- })
-# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[memory]
'''
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__;
'''
# ---
@@ -950,934 +452,16 @@
'type': 'object',
})
# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres]
- '''
- 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].1
- dict({
- '$defs': dict({
- 'InnerObject': dict({
- 'properties': dict({
- 'yo': dict({
- 'title': 'Yo',
- 'type': 'integer',
- }),
- }),
- 'required': list([
- 'yo',
- ]),
- 'title': 'InnerObject',
- 'type': 'object',
- }),
- }),
- 'properties': dict({
- 'inner': dict({
- '$ref': '#/$defs/InnerObject',
- }),
- 'query': dict({
- 'title': 'Query',
- 'type': 'string',
- }),
- }),
- 'required': list([
- 'query',
- 'inner',
- ]),
- 'title': 'Input',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres].2
- dict({
- 'properties': dict({
- 'answer': dict({
- 'title': 'Answer',
- 'type': 'string',
- }),
- 'docs': dict({
- 'items': dict({
- 'type': 'string',
- }),
- 'title': 'Docs',
- 'type': 'array',
- }),
- }),
- 'required': list([
- 'answer',
- 'docs',
- ]),
- 'title': 'Output',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_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_pipe].1
- dict({
- '$defs': dict({
- 'InnerObject': dict({
- 'properties': dict({
- 'yo': dict({
- 'title': 'Yo',
- 'type': 'integer',
- }),
- }),
- 'required': list([
- 'yo',
- ]),
- 'title': 'InnerObject',
- 'type': 'object',
- }),
- }),
- 'properties': dict({
- 'inner': dict({
- '$ref': '#/$defs/InnerObject',
- }),
- 'query': dict({
- 'title': 'Query',
- 'type': 'string',
- }),
- }),
- 'required': list([
- 'query',
- 'inner',
- ]),
- 'title': 'Input',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_pipe].2
- dict({
- 'properties': dict({
- 'answer': dict({
- 'title': 'Answer',
- 'type': 'string',
- }),
- 'docs': dict({
- 'items': dict({
- 'type': 'string',
- }),
- 'title': 'Docs',
- 'type': 'array',
- }),
- }),
- 'required': list([
- 'answer',
- 'docs',
- ]),
- 'title': 'Output',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_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_pool].1
- dict({
- '$defs': dict({
- 'InnerObject': dict({
- 'properties': dict({
- 'yo': dict({
- 'title': 'Yo',
- 'type': 'integer',
- }),
- }),
- 'required': list([
- 'yo',
- ]),
- 'title': 'InnerObject',
- 'type': 'object',
- }),
- }),
- 'properties': dict({
- 'inner': dict({
- '$ref': '#/$defs/InnerObject',
- }),
- 'query': dict({
- 'title': 'Query',
- 'type': 'string',
- }),
- }),
- 'required': list([
- 'query',
- 'inner',
- ]),
- 'title': 'Input',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_pool].2
- dict({
- 'properties': dict({
- 'answer': dict({
- 'title': 'Answer',
- 'type': 'string',
- }),
- 'docs': dict({
- 'items': dict({
- 'type': 'string',
- }),
- 'title': 'Docs',
- 'type': 'array',
- }),
- }),
- 'required': list([
- 'answer',
- 'docs',
- ]),
- 'title': 'Output',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_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_shallow].1
- dict({
- '$defs': dict({
- 'InnerObject': dict({
- 'properties': dict({
- 'yo': dict({
- 'title': 'Yo',
- 'type': 'integer',
- }),
- }),
- 'required': list([
- 'yo',
- ]),
- 'title': 'InnerObject',
- 'type': 'object',
- }),
- }),
- 'properties': dict({
- 'inner': dict({
- '$ref': '#/$defs/InnerObject',
- }),
- 'query': dict({
- 'title': 'Query',
- 'type': 'string',
- }),
- }),
- 'required': list([
- 'query',
- 'inner',
- ]),
- 'title': 'Input',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[postgres_shallow].2
- dict({
- 'properties': dict({
- 'answer': dict({
- 'title': 'Answer',
- 'type': 'string',
- }),
- 'docs': dict({
- 'items': dict({
- 'type': 'string',
- }),
- 'title': 'Docs',
- 'type': 'array',
- }),
- }),
- 'required': list([
- 'answer',
- 'docs',
- ]),
- 'title': 'Output',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite]
- '''
- 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].1
- dict({
- '$defs': dict({
- 'InnerObject': dict({
- 'properties': dict({
- 'yo': dict({
- 'title': 'Yo',
- 'type': 'integer',
- }),
- }),
- 'required': list([
- 'yo',
- ]),
- 'title': 'InnerObject',
- 'type': 'object',
- }),
- }),
- 'properties': dict({
- 'inner': dict({
- '$ref': '#/$defs/InnerObject',
- }),
- 'query': dict({
- 'title': 'Query',
- 'type': 'string',
- }),
- }),
- 'required': list([
- 'query',
- 'inner',
- ]),
- 'title': 'Input',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite].2
- dict({
- 'properties': dict({
- 'answer': dict({
- 'title': 'Answer',
- 'type': 'string',
- }),
- 'docs': dict({
- 'items': dict({
- 'type': 'string',
- }),
- 'title': 'Docs',
- 'type': 'array',
- }),
- }),
- 'required': list([
- 'answer',
- 'docs',
- ]),
- 'title': 'Output',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite_aes]
- '''
- 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_aes].1
- dict({
- '$defs': dict({
- 'InnerObject': dict({
- 'properties': dict({
- 'yo': dict({
- 'title': 'Yo',
- 'type': 'integer',
- }),
- }),
- 'required': list([
- 'yo',
- ]),
- 'title': 'InnerObject',
- 'type': 'object',
- }),
- }),
- 'properties': dict({
- 'inner': dict({
- '$ref': '#/$defs/InnerObject',
- }),
- 'query': dict({
- 'title': 'Query',
- 'type': 'string',
- }),
- }),
- 'required': list([
- 'query',
- 'inner',
- ]),
- 'title': 'Input',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic2[sqlite_aes].2
- dict({
- 'properties': dict({
- 'answer': dict({
- 'title': 'Answer',
- 'type': 'string',
- }),
- 'docs': dict({
- 'items': dict({
- 'type': 'string',
- }),
- 'title': 'Docs',
- 'type': 'array',
- }),
- }),
- 'required': list([
- 'answer',
- 'docs',
- ]),
- 'title': 'Output',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[memory]
- '''
- 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_pydantic_input[memory].1
- dict({
- '$defs': dict({
- 'InnerObject': dict({
- 'properties': dict({
- 'yo': dict({
- 'title': 'Yo',
- 'type': 'integer',
- }),
- }),
- 'required': list([
- 'yo',
- ]),
- 'title': 'InnerObject',
- 'type': 'object',
- }),
- }),
- 'properties': dict({
- 'inner': dict({
- '$ref': '#/$defs/InnerObject',
- }),
- 'query': dict({
- 'title': 'Query',
- 'type': 'string',
- }),
- }),
- 'required': list([
- 'query',
- 'inner',
- ]),
- 'title': 'Input',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[memory].2
- dict({
- 'properties': dict({
- 'answer': dict({
- 'title': 'Answer',
- 'type': 'string',
- }),
- 'docs': dict({
- 'items': dict({
- 'type': 'string',
- }),
- 'title': 'Docs',
- 'type': 'array',
- }),
- }),
- 'required': list([
- 'answer',
- 'docs',
- ]),
- 'title': 'Output',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[postgres]
- '''
- 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_pydantic_input[postgres].1
- dict({
- '$defs': dict({
- 'InnerObject': dict({
- 'properties': dict({
- 'yo': dict({
- 'title': 'Yo',
- 'type': 'integer',
- }),
- }),
- 'required': list([
- 'yo',
- ]),
- 'title': 'InnerObject',
- 'type': 'object',
- }),
- }),
- 'properties': dict({
- 'inner': dict({
- '$ref': '#/$defs/InnerObject',
- }),
- 'query': dict({
- 'title': 'Query',
- 'type': 'string',
- }),
- }),
- 'required': list([
- 'query',
- 'inner',
- ]),
- 'title': 'Input',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[postgres].2
- dict({
- 'properties': dict({
- 'answer': dict({
- 'title': 'Answer',
- 'type': 'string',
- }),
- 'docs': dict({
- 'items': dict({
- 'type': 'string',
- }),
- 'title': 'Docs',
- 'type': 'array',
- }),
- }),
- 'required': list([
- 'answer',
- 'docs',
- ]),
- 'title': 'Output',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[postgres_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_pydantic_input[postgres_pipe].1
- dict({
- '$defs': dict({
- 'InnerObject': dict({
- 'properties': dict({
- 'yo': dict({
- 'title': 'Yo',
- 'type': 'integer',
- }),
- }),
- 'required': list([
- 'yo',
- ]),
- 'title': 'InnerObject',
- 'type': 'object',
- }),
- }),
- 'properties': dict({
- 'inner': dict({
- '$ref': '#/$defs/InnerObject',
- }),
- 'query': dict({
- 'title': 'Query',
- 'type': 'string',
- }),
- }),
- 'required': list([
- 'query',
- 'inner',
- ]),
- 'title': 'Input',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[postgres_pipe].2
- dict({
- 'properties': dict({
- 'answer': dict({
- 'title': 'Answer',
- 'type': 'string',
- }),
- 'docs': dict({
- 'items': dict({
- 'type': 'string',
- }),
- 'title': 'Docs',
- 'type': 'array',
- }),
- }),
- 'required': list([
- 'answer',
- 'docs',
- ]),
- 'title': 'Output',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[postgres_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_pydantic_input[postgres_pool].1
- dict({
- '$defs': dict({
- 'InnerObject': dict({
- 'properties': dict({
- 'yo': dict({
- 'title': 'Yo',
- 'type': 'integer',
- }),
- }),
- 'required': list([
- 'yo',
- ]),
- 'title': 'InnerObject',
- 'type': 'object',
- }),
- }),
- 'properties': dict({
- 'inner': dict({
- '$ref': '#/$defs/InnerObject',
- }),
- 'query': dict({
- 'title': 'Query',
- 'type': 'string',
- }),
- }),
- 'required': list([
- 'query',
- 'inner',
- ]),
- 'title': 'Input',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[postgres_pool].2
- dict({
- 'properties': dict({
- 'answer': dict({
- 'title': 'Answer',
- 'type': 'string',
- }),
- 'docs': dict({
- 'items': dict({
- 'type': 'string',
- }),
- 'title': 'Docs',
- 'type': 'array',
- }),
- }),
- 'required': list([
- 'answer',
- 'docs',
- ]),
- 'title': 'Output',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[postgres_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_pydantic_input[postgres_shallow].1
- dict({
- '$defs': dict({
- 'InnerObject': dict({
- 'properties': dict({
- 'yo': dict({
- 'title': 'Yo',
- 'type': 'integer',
- }),
- }),
- 'required': list([
- 'yo',
- ]),
- 'title': 'InnerObject',
- 'type': 'object',
- }),
- }),
- 'properties': dict({
- 'inner': dict({
- '$ref': '#/$defs/InnerObject',
- }),
- 'query': dict({
- 'title': 'Query',
- 'type': 'string',
- }),
- }),
- 'required': list([
- 'query',
- 'inner',
- ]),
- 'title': 'Input',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[postgres_shallow].2
- dict({
- 'properties': dict({
- 'answer': dict({
- 'title': 'Answer',
- 'type': 'string',
- }),
- 'docs': dict({
- 'items': dict({
- 'type': 'string',
- }),
- 'title': 'Docs',
- 'type': 'array',
- }),
- }),
- 'required': list([
- 'answer',
- 'docs',
- ]),
- 'title': 'Output',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[sqlite]
- '''
- 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_pydantic_input[sqlite].1
- dict({
- '$defs': dict({
- 'InnerObject': dict({
- 'properties': dict({
- 'yo': dict({
- 'title': 'Yo',
- 'type': 'integer',
- }),
- }),
- 'required': list([
- 'yo',
- ]),
- 'title': 'InnerObject',
- 'type': 'object',
- }),
- }),
- 'properties': dict({
- 'inner': dict({
- '$ref': '#/$defs/InnerObject',
- }),
- 'query': dict({
- 'title': 'Query',
- 'type': 'string',
- }),
- }),
- 'required': list([
- 'query',
- 'inner',
- ]),
- 'title': 'Input',
- 'type': 'object',
- })
-# ---
-# name: test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic_input[sqlite].2
- dict({
- 'properties': dict({
- 'answer': dict({
- 'title': 'Answer',
- 'type': 'string',
- }),
- 'docs': dict({
- 'items': dict({
- 'type': 'string',
- }),
- 'title': 'Docs',
- 'type': 'array',
- }),
- }),
- 'required': list([
- 'answer',
- 'docs',
- ]),
- 'title': 'Output',
- 'type': 'object',
- })
-# ---
# name: test_in_one_fan_out_state_graph_waiting_edge_via_branch[memory]
'''
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_via_branch[postgres]
- '''
- 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_via_branch[postgres_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_via_branch[postgres_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_via_branch[postgres_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_via_branch[sqlite]
- '''
- 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_via_branch[sqlite_aes]
- '''
- 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;
'''
# ---
@@ -1918,14 +502,12 @@
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__([__start__
]):::first
+ inner(inner)
side(side)
__end__([__end__
]):::last
- __start__ --> inner_up;
- inner_up --> side;
+ __start__ --> inner;
+ inner --> side;
side --> __end__;
- subgraph inner
- inner_up(up)
- end
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
@@ -1936,56 +518,23 @@
dict({
'edges': list([
dict({
- 'conditional': True,
- 'source': 'tool_two:__start__',
- 'target': 'tool_two:tool_two_slow',
- }),
- dict({
- 'source': 'tool_two:tool_two_slow',
- 'target': 'tool_two:__end__',
- }),
- dict({
- 'conditional': True,
- 'source': 'tool_two:__start__',
- 'target': 'tool_two:tool_two_fast',
- }),
- dict({
- 'source': 'tool_two:tool_two_fast',
- 'target': 'tool_two:__end__',
- }),
- dict({
- 'conditional': True,
'source': '__start__',
- 'target': 'tool_one',
- }),
- dict({
- 'source': 'tool_one',
- 'target': '__end__',
- }),
- dict({
- 'conditional': True,
- 'source': '__start__',
- 'target': 'tool_two:__start__',
- }),
- dict({
- 'source': 'tool_two:__end__',
- 'target': '__end__',
- }),
- dict({
- 'conditional': True,
- 'source': '__start__',
- 'target': 'tool_three',
- }),
- dict({
- 'source': 'tool_three',
'target': '__end__',
}),
]),
'nodes': list([
dict({
- 'data': '__start__',
+ 'data': dict({
+ 'id': list([
+ 'langchain',
+ 'schema',
+ 'runnable',
+ 'RunnablePassthrough',
+ ]),
+ 'name': '__start__',
+ }),
'id': '__start__',
- 'type': 'schema',
+ 'type': 'runnable',
}),
dict({
'data': dict({
@@ -2000,42 +549,19 @@
'id': 'tool_one',
'type': 'runnable',
}),
- dict({
- 'data': 'tool_two:__start__',
- 'id': 'tool_two:__start__',
- 'type': 'schema',
- }),
dict({
'data': dict({
'id': list([
'langgraph',
- 'utils',
- 'runnable',
- 'RunnableCallable',
+ 'graph',
+ 'state',
+ 'CompiledStateGraph',
]),
- 'name': 'tool_two:tool_two_slow',
+ 'name': 'tool_two',
}),
- 'id': 'tool_two:tool_two_slow',
+ 'id': 'tool_two',
'type': 'runnable',
}),
- dict({
- 'data': dict({
- 'id': list([
- 'langgraph',
- 'utils',
- 'runnable',
- 'RunnableCallable',
- ]),
- 'name': 'tool_two:tool_two_fast',
- }),
- 'id': 'tool_two:tool_two_fast',
- 'type': 'runnable',
- }),
- dict({
- 'data': 'tool_two:__end__',
- 'id': 'tool_two:__end__',
- 'type': 'schema',
- }),
dict({
'data': dict({
'id': list([
@@ -2061,26 +587,12 @@
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
- __start__([__start__
]):::first
+ __start__(__start__
)
tool_one(tool_one)
+ tool_two(tool_two)
tool_three(tool_three)
- __end__([__end__
]):::last
- __start__ -.-> tool_one;
- tool_one --> __end__;
- __start__ -.-> tool_two___start__;
- tool_two___end__ --> __end__;
- __start__ -.-> tool_three;
- tool_three --> __end__;
- subgraph tool_two
- tool_two___start__(__start__
)
- tool_two_tool_two_slow(tool_two_slow)
- tool_two_tool_two_fast(tool_two_fast)
- tool_two___end__(__end__
)
- tool_two___start__ -.-> tool_two_tool_two_slow;
- tool_two_tool_two_slow --> tool_two___end__;
- tool_two___start__ -.-> tool_two_tool_two_fast;
- tool_two_tool_two_fast --> tool_two___end__;
- end
+ __end__(__end__
)
+ __start__ --> __end__;
classDef default fill:#f2f0ff,line-height:1.2
classDef first fill-opacity:0
classDef last fill:#bfb6fc
@@ -2107,11 +619,12 @@
'''
graph TD;
__start__ --> up;
- down --> __end__;
- side --> down;
- up --> down;
up --> other;
up --> side;
+ side --> down;
+ up --> down;
+ other --> __end__;
+ down --> __end__;
'''
# ---
diff --git a/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr b/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr
index 69fdad494..59d48f864 100644
--- a/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr
+++ b/libs/langgraph/tests/__snapshots__/test_pregel_async.ambr
@@ -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__([__start__
]):::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__([__start__
]):::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__([__start__
]):::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__([__start__
]):::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__([__start__
]):::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
-
- '''
-# ---
diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
index 5bac66252..2a59b06f3 100644
--- a/libs/langgraph/tests/test_pregel.py
+++ b/libs/langgraph/tests/test_pregel.py
@@ -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
diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py
index 987732b90..0ce67dc10 100644
--- a/libs/langgraph/tests/test_pregel_async.py
+++ b/libs/langgraph/tests/test_pregel_async.py
@@ -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__([__start__
]):::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
diff --git a/libs/langgraph/tests/test_pydantic.py b/libs/langgraph/tests/test_pydantic.py
index f1a350033..ad94724b5 100644
--- a/libs/langgraph/tests/test_pydantic.py
+++ b/libs/langgraph/tests/test_pydantic.py
@@ -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!"