Various improvements

Author: dmontagu

pydantic_graph/pydantic_graph/beta/decision.py

@@ -9,7 +9,7 @@
 
 from collections.abc import Callable, Iterable, Sequence
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Final, Generic
+from typing import TYPE_CHECKING, Any, Generic
 
 from typing_extensions import Never, Self, TypeVar
 
@@ -123,40 +123,29 @@ class DecisionBranch(Generic[SourceT]):
 """Type variable for transformed output."""
 
 
-@dataclass(kw_only=True)
+@dataclass(kw_only=True, frozen=True)
 class DecisionBranchBuilder(Generic[StateT, DepsT, OutputT, SourceT, HandledT]):
     """Builder for constructing decision branches with fluent API.
 
     This builder provides methods to configure branches with destinations,
     forks, and transformations in a type-safe manner.
+
+    Instances of this class should be created using [`GraphBuilder.match`][pydantic_graph.beta.graph_builder.GraphBuilder],
+    not created directly.
     """
 
-    # The use of `Final` on these attributes is necessary for them to be treated as read-only for purposes
-    # of variance-inference. This could be done with `frozen` but that
-    decision: Final[Decision[StateT, DepsT, HandledT]]
+    _decision: Decision[StateT, DepsT, HandledT]
     """The parent decision node."""
 
-    source: Final[TypeOrTypeExpression[SourceT]]
+    _source: TypeOrTypeExpression[SourceT]
     """The expected source type for this branch."""
 
-    matches: Final[Callable[[Any], bool] | None]
+    _matches: Callable[[Any], bool] | None
     """Optional matching predicate."""
 
-    path_builder: Final[PathBuilder[StateT, DepsT, OutputT]]
+    _path_builder: PathBuilder[StateT, DepsT, OutputT]
     """Builder for the execution path."""
 
-    @property
-    def last_fork_id(self) -> ForkID | None:
-        """Get the ID of the last fork in the path.
-
-        Returns:
-            The fork ID if a fork exists, None otherwise.
-        """
-        last_fork = self.path_builder.last_fork
-        if last_fork is None:
-            return None
-        return last_fork.fork_id
-
     def to(
         self,
         destination: DestinationNode[StateT, DepsT, OutputT],
@@ -173,25 +162,25 @@ def to(
             A completed DecisionBranch with the specified destinations.
         """
         return DecisionBranch(
-            source=self.source, matches=self.matches, path=self.path_builder.to(destination, *extra_destinations)
+            source=self._source, matches=self._matches, path=self._path_builder.to(destination, *extra_destinations)
         )
 
-    def fork(
+    def broadcast(
         self,
         get_forks: Callable[[Self], Sequence[DecisionBranch[SourceT]]],
         /,
     ) -> DecisionBranch[SourceT]:
-        """Create a fork in the execution path.
+        """Create a broadcast fork in the execution path.
 
         Args:
-            get_forks: Function that generates forked decision branches.
+            get_forks: Function (typically a lambda) that returns the broadcast forks downstream of this decision branch.
 
         Returns:
-            A completed DecisionBranch with forked execution paths.
+            A completed DecisionBranch with broadcast-forked execution paths.
         """
         fork_decision_branches = get_forks(self)
         new_paths = [b.path for b in fork_decision_branches]
-        return DecisionBranch(source=self.source, matches=self.matches, path=self.path_builder.fork(new_paths))
+        return DecisionBranch(source=self._source, matches=self._matches, path=self._path_builder.broadcast(new_paths))
 
     def transform(
         self, func: TransformFunction[StateT, DepsT, OutputT, NewOutputT], /
@@ -205,10 +194,10 @@ def transform(
             A new DecisionBranchBuilder where the provided transform is applied prior to generating the final output.
         """
         return DecisionBranchBuilder(
-            decision=self.decision,
-            source=self.source,
-            matches=self.matches,
-            path_builder=self.path_builder.transform(func),
+            _decision=self._decision,
+            _source=self._source,
+            _matches=self._matches,
+            _path_builder=self._path_builder.transform(func),
         )
 
     def map(
@@ -230,10 +219,10 @@ def map(
             A new DecisionBranchBuilder where mapping is performed prior to generating the final output.
         """
         return DecisionBranchBuilder(
-            decision=self.decision,
-            source=self.source,
-            matches=self.matches,
-            path_builder=self.path_builder.map(fork_id=fork_id, downstream_join_id=downstream_join_id),
+            _decision=self._decision,
+            _source=self._source,
+            _matches=self._matches,
+            _path_builder=self._path_builder.map(fork_id=fork_id, downstream_join_id=downstream_join_id),
         )
 
     def label(self, label: str) -> DecisionBranchBuilder[StateT, DepsT, OutputT, SourceT, HandledT]:
@@ -248,8 +237,8 @@ def label(self, label: str) -> DecisionBranchBuilder[StateT, DepsT, OutputT, Sou
             A new DecisionBranchBuilder where the label has been applied at the end of the current path being built.
         """
         return DecisionBranchBuilder(
-            decision=self.decision,
-            source=self.source,
-            matches=self.matches,
-            path_builder=self.path_builder.label(label),
+            _decision=self._decision,
+            _source=self._source,
+            _matches=self._matches,
+            _path_builder=self._path_builder.label(label),
         )

pydantic_graph/pydantic_graph/beta/graph_builder.py

@@ -8,9 +8,10 @@
 from __future__ import annotations
 
 import inspect
-from collections import defaultdict
+from collections import Counter, defaultdict
 from collections.abc import Callable, Iterable
-from dataclasses import dataclass
+from copy import deepcopy
+from dataclasses import dataclass, replace
 from types import NoneType
 from typing import Any, Generic, cast, get_origin, get_type_hints, overload
 
@@ -20,7 +21,7 @@
 from pydantic_graph._utils import UNSET, Unset
 from pydantic_graph.beta.decision import Decision, DecisionBranch, DecisionBranchBuilder
 from pydantic_graph.beta.graph import Graph
-from pydantic_graph.beta.id_types import ForkID, JoinID, NodeID
+from pydantic_graph.beta.id_types import ForkID, JoinID, NodeID, is_placeholder_node_id
 from pydantic_graph.beta.join import Join, JoinNode, ReducerFunction
 from pydantic_graph.beta.node import (
     EndNode,
@@ -59,6 +60,7 @@
 T = TypeVar('T', infer_variance=True)
 
 
+# TODO: Make this kw-only and drop init=False..?
 @dataclass(init=False)
 class GraphBuilder(Generic[StateT, DepsT, GraphInputT, GraphOutputT]):
     """A builder for constructing executable graph definitions.
@@ -440,7 +442,9 @@ def match(
         node_id = NodeID(self._get_new_decision_id())
         decision = Decision[StateT, DepsT, Never](node_id, branches=[], note=None)
         new_path_builder = PathBuilder[StateT, DepsT, SourceT](working_items=[])
-        return DecisionBranchBuilder(decision=decision, source=source, matches=matches, path_builder=new_path_builder)
+        return DecisionBranchBuilder(
+            _decision=decision, _source=source, _matches=matches, _path_builder=new_path_builder
+        )
 
     def match_node(
         self,
@@ -663,8 +667,11 @@ def build(self) -> Graph[StateT, DepsT, GraphInputT, GraphOutputT]:
         # TODO(P2): Allow the user to specify the parent forks; only infer them if _not_ specified
         # TODO(P2): Verify that any user-specified parent forks are _actually_ valid parent forks, and if not, generate a helpful error message
         # TODO(P3): Consider doing a deepcopy here to prevent modifications to the underlying nodes and edges
-        nodes = self._nodes
-        edges_by_source = self._edges_by_source
+
+        nodes = deepcopy(self._nodes)
+        edges_by_source = deepcopy(self._edges_by_source)
+
+        nodes, edges_by_source = _replace_placeholder_node_ids(nodes, edges_by_source)
         nodes, edges_by_source = _flatten_paths(nodes, edges_by_source)
         nodes, edges_by_source = _normalize_forks(nodes, edges_by_source)
         parent_forks = _collect_dominating_forks(nodes, edges_by_source)
@@ -843,3 +850,81 @@ def _handle_path(path: Path, last_source_id: NodeID):
         dominating_forks[join_id] = dominating_fork
 
     return dominating_forks
+
+
+def _replace_placeholder_node_ids(nodes: dict[NodeID, AnyNode], edges_by_source: dict[NodeID, list[Path]]):
+    node_id_remapping = _build_placeholder_node_id_remapping(nodes)
+    replaced_nodes = {
+        node_id_remapping.get(name, name): _update_node_with_id_remapping(node, node_id_remapping)
+        for name, node in nodes.items()
+    }
+    replaced_edges_by_source = {
+        node_id_remapping.get(source, source): [_update_path_with_id_remapping(p, node_id_remapping) for p in paths]
+        for source, paths in edges_by_source.items()
+    }
+    return replaced_nodes, replaced_edges_by_source
+
+
+def _build_placeholder_node_id_remapping(nodes: dict[NodeID, AnyNode]) -> dict[NodeID, NodeID]:
+    """The determinism of the generated remapping here is dependent on the determinism of the ordering of the `nodes` dict.
+
+    Note: If we want to generate more interesting names, we could try to make use of information about the edges
+    into/out of the relevant nodes. I'm not sure if there's a good use case for that though so I didn't bother for now.
+    """
+    counter = Counter[str]()
+    remapping: dict[NodeID, NodeID] = {}
+    for node_id, node in nodes.items():
+        if not is_placeholder_node_id(node_id):
+            continue
+        label = type(node).__name__.lower()
+        counter[label] = count = counter[label] + 1
+        remapping[node_id] = NodeID(f'{label}_{count}')
+    return remapping
+
+
+def _update_node_with_id_remapping(node: AnyNode, node_id_remapping: dict[NodeID, NodeID]) -> AnyNode:
+    if isinstance(node, Step):
+        # Even though steps are frozen, we use object.__setattr__ to overrule that and change the id value to make it
+        # work with NodeStep.
+        # Note: we have already deepcopied the inputs to this function so it should be okay to make mutations,
+        # this could change if we change the code surrounding the code paths leading to this function call though.
+        object.__setattr__(node, 'id', node_id_remapping.get(node.id, node.id))
+    elif isinstance(node, Join):
+        node = replace(node, id=JoinID(node_id_remapping.get(node.id, node.id)))
+    elif isinstance(node, Fork):
+        node = replace(node, id=ForkID(node_id_remapping.get(node.id, node.id)))
+    elif isinstance(node, Decision):
+        node = replace(
+            node,
+            id=node_id_remapping.get(node.id, node.id),
+            branches=[
+                replace(branch, path=_update_path_with_id_remapping(branch.path, node_id_remapping))
+                for branch in node.branches
+            ],
+        )
+    return node
+
+
+def _update_path_with_id_remapping(path: Path, node_id_remapping: dict[NodeID, NodeID]) -> Path:
+    path = replace(path)  # prevent mutating the input; not technically necessary but could make debugging easier later
+    for i, item in enumerate(path.items):
+        if isinstance(item, MapMarker):
+            downstream_join_id = item.downstream_join_id
+            if downstream_join_id is not None:
+                downstream_join_id = JoinID(node_id_remapping.get(downstream_join_id, downstream_join_id))
+            path.items[i] = replace(
+                item,
+                fork_id=ForkID(node_id_remapping.get(item.fork_id, item.fork_id)),
+                downstream_join_id=downstream_join_id,
+            )
+        elif isinstance(item, BroadcastMarker):
+            path.items[i] = replace(
+                item,
+                fork_id=ForkID(node_id_remapping.get(item.fork_id, item.fork_id)),
+                paths=[_update_path_with_id_remapping(p, node_id_remapping) for p in item.paths],
+            )
+        elif isinstance(item, DestinationMarker):
+            path.items[i] = replace(
+                item, destination_id=node_id_remapping.get(item.destination_id, item.destination_id)
+            )
+    return path

pydantic_graph/pydantic_graph/beta/id_types.py

@@ -6,6 +6,7 @@
 
 from __future__ import annotations
 
+import uuid
 from dataclasses import dataclass
 from typing import NewType
 
@@ -54,3 +55,24 @@ class ForkStackItem:
 The fork stack tracks the complete path through nested parallel executions,
 allowing the system to coordinate and join parallel branches correctly.
 """
+
+
+def generate_placeholder_node_id(debug_label: str) -> str:
+    """Generate a placeholder node ID, to be replaced during graph building."""
+    return f'{_NODE_ID_PLACEHOLDER_PREFIX}:{debug_label}:{uuid.uuid4()}'
+
+
+def is_placeholder_node_id(node_id: NodeID) -> bool:
+    """Returns whether a given NodeID is a placeholder node ID which should be replaced during graph building."""
+    return node_id.startswith(_NODE_ID_PLACEHOLDER_PREFIX)
+
+
+_NODE_ID_PLACEHOLDER_PREFIX = '__placeholder__:'
+"""
+When Node IDs are required but not specified when building a graph, we generate placeholder values
+using this prefix followed by a random string.
+
+During graph building, we replace these with simpler and deterministically-selected values.
+This ensures that the node IDs are stable when rebuilding the graph, and makes the generated mermaid diagrams etc.
+easier to read.
+"""

pydantic_graph/pydantic_graph/beta/mermaid.py

@@ -131,7 +131,8 @@ def render(
         if direction is not None:
             lines.append(f'  direction {direction}')
 
-        for node in self.nodes:
+        nodes, edges = _topological_sort(self.nodes, self.edges)
+        for node in nodes:
             # List all nodes in order they were created
             node_lines: list[str] = []
             if node.kind == 'start' or node.kind == 'end':
@@ -156,7 +157,7 @@ def render(
 
         lines.append('')
 
-        for edge in self.edges:
+        for edge in edges:
             # Use special [*] syntax for start/end nodes
             render_start_id = '[*]' if edge.start_id == StartNode.id else edge.start_id
             render_end_id = '[*]' if edge.end_id == EndNode.id else edge.end_id
@@ -167,3 +168,48 @@ def render(
             # TODO(P3): Support node notes/highlighting
 
         return '\n'.join(lines)
+
+
+def _topological_sort(
+    nodes: list[MermaidNode], edges: list[MermaidEdge]
+) -> tuple[list[MermaidNode], list[MermaidEdge]]:
+    """Sort nodes and edges in a logical topological order.
+
+    Uses BFS from the start node to assign depths, then sorts:
+    - Nodes by their distance from start
+    - Edges by the distance of their source and target nodes
+    """
+    # Build adjacency list for BFS
+    adjacency: dict[str, list[str]] = defaultdict(list)
+    for edge in edges:
+        adjacency[edge.start_id].append(edge.end_id)
+
+    # BFS to assign depth to each node (distance from start)
+    depths: dict[str, int] = {}
+    queue: list[tuple[str, int]] = [(StartNode.id, 0)]
+    depths[StartNode.id] = 0
+
+    while queue:
+        node_id, depth = queue.pop(0)
+        for next_id in adjacency[node_id]:
+            if next_id not in depths:
+                depths[next_id] = depth + 1
+                queue.append((next_id, depth + 1))
+
+    # Sort nodes by depth (distance from start), then by id for stability
+    # Nodes not reachable from start get infinity depth (sorted to end)
+    sorted_nodes = sorted(nodes, key=lambda n: (depths.get(n.id, float('inf')), n.id))
+
+    # Sort edges by source depth, then target depth
+    # This ensures edges closer to start come first, edges closer to end come last
+    sorted_edges = sorted(
+        edges,
+        key=lambda e: (
+            depths.get(e.start_id, float('inf')),
+            depths.get(e.end_id, float('inf')),
+            e.start_id,
+            e.end_id,
+        ),
+    )
+
+    return sorted_nodes, sorted_edges