pydantic-graph 0.7.5__py3-none-any.whl → 1.24.0__py3-none-any.whl

pydantic_graph/beta/graph_builder.py

@@ -0,0 +1,1053 @@
+"""Graph builder for constructing executable graph definitions.
+
+This module provides the GraphBuilder class and related utilities for
+constructing typed, executable graph definitions with steps, joins,
+decisions, and edge routing.
+"""
+
+from __future__ import annotations
+
+import inspect
+from collections import Counter, defaultdict
+from collections.abc import AsyncIterable, Callable, Iterable
+from dataclasses import dataclass, replace
+from types import NoneType
+from typing import Any, Generic, Literal, cast, get_origin, get_type_hints, overload
+
+from typing_extensions import Never, TypeAliasType, TypeVar
+
+from pydantic_graph import _utils, exceptions
+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, generate_placeholder_node_id, replace_placeholder_id
+from pydantic_graph.beta.join import Join, JoinNode, ReducerFunction
+from pydantic_graph.beta.mermaid import build_mermaid_graph
+from pydantic_graph.beta.node import (
+    EndNode,
+    Fork,
+    StartNode,
+)
+from pydantic_graph.beta.node_types import (
+    AnyDestinationNode,
+    AnyNode,
+    DestinationNode,
+    SourceNode,
+)
+from pydantic_graph.beta.parent_forks import ParentFork, ParentForkFinder
+from pydantic_graph.beta.paths import (
+    BroadcastMarker,
+    DestinationMarker,
+    EdgePath,
+    EdgePathBuilder,
+    MapMarker,
+    Path,
+    PathBuilder,
+)
+from pydantic_graph.beta.step import NodeStep, Step, StepContext, StepFunction, StepNode, StreamFunction
+from pydantic_graph.beta.util import TypeOrTypeExpression, get_callable_name, unpack_type_expression
+from pydantic_graph.exceptions import GraphBuildingError, GraphValidationError
+from pydantic_graph.nodes import BaseNode, End
+
+StateT = TypeVar('StateT', infer_variance=True)
+DepsT = TypeVar('DepsT', infer_variance=True)
+InputT = TypeVar('InputT', infer_variance=True)
+OutputT = TypeVar('OutputT', infer_variance=True)
+SourceT = TypeVar('SourceT', infer_variance=True)
+SourceNodeT = TypeVar('SourceNodeT', bound=BaseNode[Any, Any, Any], infer_variance=True)
+SourceOutputT = TypeVar('SourceOutputT', infer_variance=True)
+GraphInputT = TypeVar('GraphInputT', infer_variance=True)
+GraphOutputT = TypeVar('GraphOutputT', infer_variance=True)
+T = TypeVar('T', infer_variance=True)
+
+
+@dataclass(init=False)
+class GraphBuilder(Generic[StateT, DepsT, GraphInputT, GraphOutputT]):
+    """A builder for constructing executable graph definitions.
+
+    GraphBuilder provides a fluent interface for defining nodes, edges, and
+    routing in a graph workflow. It supports typed state, dependencies, and
+    input/output validation.
+
+    Type Parameters:
+        StateT: The type of the graph state
+        DepsT: The type of the dependencies
+        GraphInputT: The type of the graph input data
+        GraphOutputT: The type of the graph output data
+    """
+
+    name: str | None
+    """Optional name for the graph, if not provided the name will be inferred from the calling frame on the first call to a graph method."""
+
+    state_type: TypeOrTypeExpression[StateT]
+    """The type of the graph state."""
+
+    deps_type: TypeOrTypeExpression[DepsT]
+    """The type of the dependencies."""
+
+    input_type: TypeOrTypeExpression[GraphInputT]
+    """The type of the graph input data."""
+
+    output_type: TypeOrTypeExpression[GraphOutputT]
+    """The type of the graph output data."""
+
+    auto_instrument: bool
+    """Whether to automatically create instrumentation spans."""
+
+    _nodes: dict[NodeID, AnyNode]
+    """Internal storage for nodes in the graph."""
+
+    _edges_by_source: dict[NodeID, list[Path]]
+    """Internal storage for edges by source node."""
+
+    _decision_index: int
+    """Counter for generating unique decision node IDs."""
+
+    Source = TypeAliasType('Source', SourceNode[StateT, DepsT, OutputT], type_params=(OutputT,))
+    Destination = TypeAliasType('Destination', DestinationNode[StateT, DepsT, InputT], type_params=(InputT,))
+
+    def __init__(
+        self,
+        *,
+        name: str | None = None,
+        state_type: TypeOrTypeExpression[StateT] = NoneType,
+        deps_type: TypeOrTypeExpression[DepsT] = NoneType,
+        input_type: TypeOrTypeExpression[GraphInputT] = NoneType,
+        output_type: TypeOrTypeExpression[GraphOutputT] = NoneType,
+        auto_instrument: bool = True,
+    ):
+        """Initialize a graph builder.
+
+        Args:
+            name: Optional name for the graph, if not provided the name will be inferred from the calling frame on the first call to a graph method.
+            state_type: The type of the graph state
+            deps_type: The type of the dependencies
+            input_type: The type of the graph input data
+            output_type: The type of the graph output data
+            auto_instrument: Whether to automatically create instrumentation spans
+        """
+        self.name = name
+
+        self.state_type = state_type
+        self.deps_type = deps_type
+        self.input_type = input_type
+        self.output_type = output_type
+
+        self.auto_instrument = auto_instrument
+
+        self._nodes = {}
+        self._edges_by_source = defaultdict(list)
+        self._decision_index = 1
+
+        self._start_node = StartNode[GraphInputT]()
+        self._end_node = EndNode[GraphOutputT]()
+
+    # Node building
+    @property
+    def start_node(self) -> StartNode[GraphInputT]:
+        """Get the start node for the graph.
+
+        Returns:
+            The start node that receives the initial graph input
+        """
+        return self._start_node
+
+    @property
+    def end_node(self) -> EndNode[GraphOutputT]:
+        """Get the end node for the graph.
+
+        Returns:
+            The end node that produces the final graph output
+        """
+        return self._end_node
+
+    @overload
+    def step(
+        self,
+        *,
+        node_id: str | None = None,
+        label: str | None = None,
+    ) -> Callable[[StepFunction[StateT, DepsT, InputT, OutputT]], Step[StateT, DepsT, InputT, OutputT]]: ...
+    @overload
+    def step(
+        self,
+        call: StepFunction[StateT, DepsT, InputT, OutputT],
+        *,
+        node_id: str | None = None,
+        label: str | None = None,
+    ) -> Step[StateT, DepsT, InputT, OutputT]: ...
+    def step(
+        self,
+        call: StepFunction[StateT, DepsT, InputT, OutputT] | None = None,
+        *,
+        node_id: str | None = None,
+        label: str | None = None,
+    ) -> (
+        Step[StateT, DepsT, InputT, OutputT]
+        | Callable[[StepFunction[StateT, DepsT, InputT, OutputT]], Step[StateT, DepsT, InputT, OutputT]]
+    ):
+        """Create a step from a step function.
+
+        This method can be used as a decorator or called directly to create
+        a step node from an async function.
+
+        Args:
+            call: The step function to wrap
+            node_id: Optional ID for the node
+            label: Optional human-readable label
+
+        Returns:
+            Either a Step instance or a decorator function
+        """
+        if call is None:
+
+            def decorator(
+                func: StepFunction[StateT, DepsT, InputT, OutputT],
+            ) -> Step[StateT, DepsT, InputT, OutputT]:
+                return self.step(call=func, node_id=node_id, label=label)
+
+            return decorator
+
+        node_id = node_id or get_callable_name(call)
+
+        step = Step[StateT, DepsT, InputT, OutputT](id=NodeID(node_id), call=call, label=label)
+
+        return step
+
+    @overload
+    def stream(
+        self,
+        *,
+        node_id: str | None = None,
+        label: str | None = None,
+    ) -> Callable[
+        [StreamFunction[StateT, DepsT, InputT, OutputT]], Step[StateT, DepsT, InputT, AsyncIterable[OutputT]]
+    ]: ...
+    @overload
+    def stream(
+        self,
+        call: StreamFunction[StateT, DepsT, InputT, OutputT],
+        *,
+        node_id: str | None = None,
+        label: str | None = None,
+    ) -> Step[StateT, DepsT, InputT, AsyncIterable[OutputT]]: ...
+    @overload
+    def stream(
+        self,
+        call: StreamFunction[StateT, DepsT, InputT, OutputT] | None = None,
+        *,
+        node_id: str | None = None,
+        label: str | None = None,
+    ) -> (
+        Step[StateT, DepsT, InputT, AsyncIterable[OutputT]]
+        | Callable[
+            [StreamFunction[StateT, DepsT, InputT, OutputT]],
+            Step[StateT, DepsT, InputT, AsyncIterable[OutputT]],
+        ]
+    ): ...
+    def stream(
+        self,
+        call: StreamFunction[StateT, DepsT, InputT, OutputT] | None = None,
+        *,
+        node_id: str | None = None,
+        label: str | None = None,
+    ) -> (
+        Step[StateT, DepsT, InputT, AsyncIterable[OutputT]]
+        | Callable[
+            [StreamFunction[StateT, DepsT, InputT, OutputT]],
+            Step[StateT, DepsT, InputT, AsyncIterable[OutputT]],
+        ]
+    ):
+        """Create a step from an async iterator (which functions like a "stream").
+
+        This method can be used as a decorator or called directly to create
+        a step node from an async function.
+
+        Args:
+            call: The step function to wrap
+            node_id: Optional ID for the node
+            label: Optional human-readable label
+
+        Returns:
+            Either a Step instance or a decorator function
+        """
+        if call is None:
+
+            def decorator(
+                func: StreamFunction[StateT, DepsT, InputT, OutputT],
+            ) -> Step[StateT, DepsT, InputT, AsyncIterable[OutputT]]:
+                return self.stream(call=func, node_id=node_id, label=label)
+
+            return decorator
+
+        # We need to wrap the call so that we can call `await` even though the result is an async iterator
+        async def wrapper(ctx: StepContext[StateT, DepsT, InputT]):
+            return call(ctx)
+
+        return self.step(call=wrapper, node_id=node_id, label=label)
+
+    @overload
+    def join(
+        self,
+        reducer: ReducerFunction[StateT, DepsT, InputT, OutputT],
+        *,
+        initial: OutputT,
+        node_id: str | None = None,
+        parent_fork_id: str | None = None,
+        preferred_parent_fork: Literal['farthest', 'closest'] = 'farthest',
+    ) -> Join[StateT, DepsT, InputT, OutputT]: ...
+    @overload
+    def join(
+        self,
+        reducer: ReducerFunction[StateT, DepsT, InputT, OutputT],
+        *,
+        initial_factory: Callable[[], OutputT],
+        node_id: str | None = None,
+        parent_fork_id: str | None = None,
+        preferred_parent_fork: Literal['farthest', 'closest'] = 'farthest',
+    ) -> Join[StateT, DepsT, InputT, OutputT]: ...
+
+    def join(
+        self,
+        reducer: ReducerFunction[StateT, DepsT, InputT, OutputT],
+        *,
+        initial: OutputT | Unset = UNSET,
+        initial_factory: Callable[[], OutputT] | Unset = UNSET,
+        node_id: str | None = None,
+        parent_fork_id: str | None = None,
+        preferred_parent_fork: Literal['farthest', 'closest'] = 'farthest',
+    ) -> Join[StateT, DepsT, InputT, OutputT]:
+        if initial_factory is UNSET:
+            initial_factory = lambda: initial  # pyright: ignore[reportAssignmentType]  # noqa: E731
+
+        return Join[StateT, DepsT, InputT, OutputT](
+            id=JoinID(NodeID(node_id or generate_placeholder_node_id(get_callable_name(reducer)))),
+            reducer=reducer,
+            initial_factory=cast(Callable[[], OutputT], initial_factory),
+            parent_fork_id=ForkID(parent_fork_id) if parent_fork_id is not None else None,
+            preferred_parent_fork=preferred_parent_fork,
+        )
+
+    # Edge building
+    def add(self, *edges: EdgePath[StateT, DepsT]) -> None:  # noqa: C901
+        """Add one or more edge paths to the graph.
+
+        This method processes edge paths and automatically creates any necessary
+        fork nodes for broadcasts and maps.
+
+        Args:
+            *edges: The edge paths to add to the graph
+        """
+
+        def _handle_path(p: Path):
+            """Process a path and create necessary fork nodes.
+
+            Args:
+                p: The path to process
+            """
+            for item in p.items:
+                if isinstance(item, BroadcastMarker):
+                    new_node = Fork[Any, Any](id=item.fork_id, is_map=False, downstream_join_id=None)
+                    self._insert_node(new_node)
+                    for path in item.paths:
+                        _handle_path(Path(items=[*path.items]))
+                elif isinstance(item, MapMarker):
+                    new_node = Fork[Any, Any](id=item.fork_id, is_map=True, downstream_join_id=item.downstream_join_id)
+                    self._insert_node(new_node)
+                elif isinstance(item, DestinationMarker):
+                    pass
+
+        def _handle_destination_node(d: AnyDestinationNode):
+            if id(d) in destination_ids:
+                return  # prevent infinite recursion if there is a cycle of decisions
+
+            destination_ids.add(id(d))
+            destinations.append(d)
+            self._insert_node(d)
+            if isinstance(d, Decision):
+                for branch in d.branches:
+                    _handle_path(branch.path)
+                    for d2 in branch.destinations:
+                        _handle_destination_node(d2)
+
+        destination_ids = set[int]()
+        destinations: list[AnyDestinationNode] = []
+        for edge in edges:
+            for source_node in edge.sources:
+                self._insert_node(source_node)
+                self._edges_by_source[source_node.id].append(edge.path)
+            for destination_node in edge.destinations:
+                _handle_destination_node(destination_node)
+            _handle_path(edge.path)
+
+        # Automatically create edges from step function return hints including `BaseNode`s
+        for destination in destinations:
+            if not isinstance(destination, Step) or isinstance(destination, NodeStep):
+                continue
+            parent_namespace = _utils.get_parent_namespace(inspect.currentframe())
+            type_hints = get_type_hints(destination.call, localns=parent_namespace, include_extras=True)
+            try:
+                return_hint = type_hints['return']
+            except KeyError:
+                pass
+            else:
+                edge = self._edge_from_return_hint(destination, return_hint)
+                if edge is not None:
+                    self.add(edge)
+
+    def add_edge(self, source: Source[T], destination: Destination[T], *, label: str | None = None) -> None:
+        """Add a simple edge between two nodes.
+
+        Args:
+            source: The source node
+            destination: The destination node
+            label: Optional label for the edge
+        """
+        builder = self.edge_from(source)
+        if label is not None:
+            builder = builder.label(label)
+        self.add(builder.to(destination))
+
+    def add_mapping_edge(
+        self,
+        source: Source[Iterable[T]],
+        map_to: Destination[T],
+        *,
+        pre_map_label: str | None = None,
+        post_map_label: str | None = None,
+        fork_id: ForkID | None = None,
+        downstream_join_id: JoinID | None = None,
+    ) -> None:
+        """Add an edge that maps iterable data across parallel paths.
+
+        Args:
+            source: The source node that produces iterable data
+            map_to: The destination node that receives individual items
+            pre_map_label: Optional label before the map operation
+            post_map_label: Optional label after the map operation
+            fork_id: Optional ID for the fork node produced for this map operation
+            downstream_join_id: Optional ID of a join node that will always be downstream of this map.
+                Specifying this ensures correct handling if you try to map an empty iterable.
+        """
+        builder = self.edge_from(source)
+        if pre_map_label is not None:
+            builder = builder.label(pre_map_label)
+        builder = builder.map(fork_id=fork_id, downstream_join_id=downstream_join_id)
+        if post_map_label is not None:
+            builder = builder.label(post_map_label)
+        self.add(builder.to(map_to))
+
+    # TODO(DavidM): Support adding subgraphs; I think this behaves like a step with the same inputs/outputs but gets rendered as a subgraph in mermaid
+
+    def edge_from(self, *sources: Source[SourceOutputT]) -> EdgePathBuilder[StateT, DepsT, SourceOutputT]:
+        """Create an edge path builder starting from the given source nodes.
+
+        Args:
+            *sources: The source nodes to start the edge path from
+
+        Returns:
+            An EdgePathBuilder for constructing the complete edge path
+        """
+        return EdgePathBuilder[StateT, DepsT, SourceOutputT](
+            sources=sources, path_builder=PathBuilder(working_items=[])
+        )
+
+    def decision(self, *, note: str | None = None, node_id: str | None = None) -> Decision[StateT, DepsT, Never]:
+        """Create a new decision node.
+
+        Args:
+            note: Optional note to describe the decision logic
+            node_id: Optional ID for the node produced for this decision logic
+
+        Returns:
+            A new Decision node with no branches
+        """
+        return Decision(id=NodeID(node_id or generate_placeholder_node_id('decision')), branches=[], note=note)
+
+    def match(
+        self,
+        source: TypeOrTypeExpression[SourceT],
+        *,
+        matches: Callable[[Any], bool] | None = None,
+    ) -> DecisionBranchBuilder[StateT, DepsT, SourceT, SourceT, Never]:
+        """Create a decision branch matcher.
+
+        Args:
+            source: The type or type expression to match against
+            matches: Optional custom matching function
+
+        Returns:
+            A DecisionBranchBuilder for constructing the branch
+        """
+        # Note, the following node_id really is just a placeholder and shouldn't end up in the final graph
+        # This is why we don't expose a way for end users to override the value used here.
+        node_id = NodeID(generate_placeholder_node_id('match_decision'))
+        decision = Decision[StateT, DepsT, Never](id=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)
+
+    def match_node(
+        self,
+        source: type[SourceNodeT],
+        *,
+        matches: Callable[[Any], bool] | None = None,
+    ) -> DecisionBranch[SourceNodeT]:
+        """Create a decision branch for BaseNode subclasses.
+
+        This is similar to match() but specifically designed for matching
+        against BaseNode types from the v1 system.
+
+        Args:
+            source: The BaseNode subclass to match against
+            matches: Optional custom matching function
+
+        Returns:
+            A DecisionBranch for the BaseNode type
+        """
+        node = NodeStep(source)
+        path = Path(items=[DestinationMarker(node.id)])
+        return DecisionBranch(source=source, matches=matches, path=path, destinations=[node])
+
+    def node(
+        self,
+        node_type: type[BaseNode[StateT, DepsT, GraphOutputT]],
+    ) -> EdgePath[StateT, DepsT]:
+        """Create an edge path from a BaseNode class.
+
+        This method integrates v1-style BaseNode classes into the v2 graph
+        system by analyzing their type hints and creating appropriate edges.
+
+        Args:
+            node_type: The BaseNode subclass to integrate
+
+        Returns:
+            An EdgePath representing the node and its connections
+
+        Raises:
+            GraphSetupError: If the node type is missing required type hints
+        """
+        parent_namespace = _utils.get_parent_namespace(inspect.currentframe())
+        type_hints = get_type_hints(node_type.run, localns=parent_namespace, include_extras=True)
+        try:
+            return_hint = type_hints['return']
+        except KeyError as e:  # pragma: no cover
+            raise exceptions.GraphSetupError(
+                f'Node {node_type} is missing a return type hint on its `run` method'
+            ) from e
+
+        node = NodeStep(node_type)
+
+        edge = self._edge_from_return_hint(node, return_hint)
+        if not edge:  # pragma: no cover
+            raise exceptions.GraphSetupError(f'Node {node_type} is missing a return type hint on its `run` method')
+
+        return edge
+
+    # Helpers
+    def _insert_node(self, node: AnyNode) -> None:
+        """Insert a node into the graph, checking for ID conflicts.
+
+        Args:
+            node: The node to insert
+
+        Raises:
+            ValueError: If a different node with the same ID already exists
+        """
+        existing = self._nodes.get(node.id)
+        if existing is None:
+            self._nodes[node.id] = node
+        elif isinstance(existing, NodeStep) and isinstance(node, NodeStep) and existing.node_type is node.node_type:
+            pass
+        elif existing is not node:
+            raise GraphBuildingError(
+                f'All nodes must have unique node IDs. {node.id!r} was the ID for {existing} and {node}'
+            )
+
+    def _edge_from_return_hint(
+        self, node: SourceNode[StateT, DepsT, Any], return_hint: TypeOrTypeExpression[Any]
+    ) -> EdgePath[StateT, DepsT] | None:
+        """Create edges from a return type hint.
+
+        This method analyzes return type hints from step functions or node methods
+        to automatically create appropriate edges in the graph.
+
+        Args:
+            node: The source node
+            return_hint: The return type hint to analyze
+
+        Returns:
+            An EdgePath if edges can be inferred, None otherwise
+
+        Raises:
+            GraphSetupError: If the return type hint is invalid or incomplete
+        """
+        destinations: list[AnyDestinationNode] = []
+        union_args = _utils.get_union_args(return_hint)
+        for return_type in union_args:
+            return_type, annotations = _utils.unpack_annotated(return_type)
+            return_type_origin = get_origin(return_type) or return_type
+            if return_type_origin is End:
+                destinations.append(self.end_node)
+            elif return_type_origin is BaseNode:
+                raise exceptions.GraphSetupError(  # pragma: no cover
+                    f'Node {node} return type hint includes a plain `BaseNode`. '
+                    'Edge inference requires each possible returned `BaseNode` subclass to be listed explicitly.'
+                )
+            elif return_type_origin is StepNode:
+                step = cast(
+                    Step[StateT, DepsT, Any, Any] | None,
+                    next((a for a in annotations if isinstance(a, Step)), None),  # pyright: ignore[reportUnknownArgumentType]
+                )
+                if step is None:
+                    raise exceptions.GraphSetupError(  # pragma: no cover
+                        f'Node {node} return type hint includes a `StepNode` without a `Step` annotation. '
+                        'When returning `my_step.as_node()`, use `Annotated[StepNode[StateT, DepsT], my_step]` as the return type hint.'
+                    )
+                destinations.append(step)
+            elif return_type_origin is JoinNode:
+                join = cast(
+                    Join[StateT, DepsT, Any, Any] | None,
+                    next((a for a in annotations if isinstance(a, Join)), None),  # pyright: ignore[reportUnknownArgumentType]
+                )
+                if join is None:
+                    raise exceptions.GraphSetupError(  # pragma: no cover
+                        f'Node {node} return type hint includes a `JoinNode` without a `Join` annotation. '
+                        'When returning `my_join.as_node()`, use `Annotated[JoinNode[StateT, DepsT], my_join]` as the return type hint.'
+                    )
+                destinations.append(join)
+            elif inspect.isclass(return_type_origin) and issubclass(return_type_origin, BaseNode):
+                destinations.append(NodeStep(return_type))
+
+        if len(destinations) < len(union_args):
+            # Only build edges if all the return types are nodes
+            return None
+
+        edge = self.edge_from(node)
+        if len(destinations) == 1:
+            return edge.to(destinations[0])
+        else:
+            decision = self.decision()
+            for destination in destinations:
+                # We don't actually use this decision mechanism, but we need to build the edges for parent-fork finding
+                decision = decision.branch(self.match(NoneType).to(destination))
+            return edge.to(decision)
+
+    # Graph building
+    def build(self, validate_graph_structure: bool = True) -> Graph[StateT, DepsT, GraphInputT, GraphOutputT]:
+        """Build the final executable graph from the accumulated nodes and edges.
+
+        This method performs validation, normalization, and analysis of the graph
+        structure to create a complete, executable graph instance.
+
+        Args:
+            validate_graph_structure: whether to perform validation of the graph structure
+                See the docstring of _validate_graph_structure below for more details.
+
+        Returns:
+            A complete Graph instance ready for execution
+
+        Raises:
+            ValueError: If the graph structure is invalid (e.g., join without parent fork)
+        """
+        nodes = self._nodes
+        edges_by_source = 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)
+        if validate_graph_structure:
+            _validate_graph_structure(nodes, edges_by_source)
+        parent_forks = _collect_dominating_forks(nodes, edges_by_source)
+        intermediate_join_nodes = _compute_intermediate_join_nodes(nodes, parent_forks)
+
+        return Graph[StateT, DepsT, GraphInputT, GraphOutputT](
+            name=self.name,
+            state_type=unpack_type_expression(self.state_type),
+            deps_type=unpack_type_expression(self.deps_type),
+            input_type=unpack_type_expression(self.input_type),
+            output_type=unpack_type_expression(self.output_type),
+            nodes=nodes,
+            edges_by_source=edges_by_source,
+            parent_forks=parent_forks,
+            intermediate_join_nodes=intermediate_join_nodes,
+            auto_instrument=self.auto_instrument,
+        )
+
+
+def _validate_graph_structure(  # noqa: C901
+    nodes: dict[NodeID, AnyNode],
+    edges_by_source: dict[NodeID, list[Path]],
+) -> None:
+    """Validate the graph structure for common issues.
+
+    This function raises an error if any of the following criteria are not met:
+    1. There are edges from the start node
+    2. There are edges to the end node
+    3. No non-End node is a dead end (no outgoing edges)
+    4. The end node is reachable from the start node
+    5. All nodes are reachable from the start node
+
+    Note 1: Under some circumstances it may be reasonable to build a graph that violates one or more of
+    the above conditions. We may eventually add support for more granular control over validation,
+    but today, if you want to build a graph that violates any of these assumptions you need to pass
+    `validate_graph_structure=False` to the call to `GraphBuilder.build`.
+
+    Note 2: Some of the earlier items in the above list are redundant with the later items.
+    I've included the earlier items in the list as a reminder to ourselves if/when we add more granular validation
+    because you might want to check the earlier items but not the later items, as described in Note 1.
+
+    Args:
+        nodes: The nodes in the graph
+        edges_by_source: The edges by source node
+
+    Raises:
+        GraphBuildingError: If any of the aforementioned structural issues are found.
+    """
+    how_to_suppress = ' If this is intentional, you can suppress this error by passing `validate_graph_structure=False` to the call to `GraphBuilder.build`.'
+
+    # Extract all destination IDs from edges and decision branches
+    all_destinations: set[NodeID] = set()
+
+    def _collect_destinations_from_path(path: Path) -> None:
+        for item in path.items:
+            if isinstance(item, DestinationMarker):
+                all_destinations.add(item.destination_id)
+
+    for paths in edges_by_source.values():
+        for path in paths:
+            _collect_destinations_from_path(path)
+
+    # Also collect destinations from decision branches
+    for node in nodes.values():
+        if isinstance(node, Decision):
+            for branch in node.branches:
+                _collect_destinations_from_path(branch.path)
+
+    # Check 1: Check if there are edges from the start node
+    start_edges = edges_by_source.get(StartNode.id, [])
+    if not start_edges:
+        raise GraphValidationError('The graph has no edges from the start node.' + how_to_suppress)
+
+    # Check 2: Check if there are edges to the end node
+    if EndNode.id not in all_destinations:
+        raise GraphValidationError('The graph has no edges to the end node.' + how_to_suppress)
+
+    # Check 3: Find all nodes with no outgoing edges (dead ends)
+    dead_end_nodes: list[NodeID] = []
+    for node_id, node in nodes.items():
+        # Skip the end node itself
+        if isinstance(node, EndNode):
+            continue
+
+        # Check if this node has any outgoing edges
+        has_edges = node_id in edges_by_source and len(edges_by_source[node_id]) > 0
+
+        # Also check if it's a decision node with branches
+        if isinstance(node, Decision):
+            has_edges = has_edges or len(node.branches) > 0
+
+        if not has_edges:
+            dead_end_nodes.append(node_id)
+
+    if dead_end_nodes:
+        raise GraphValidationError(f'The following nodes have no outgoing edges: {dead_end_nodes}.' + how_to_suppress)
+
+    # Checks 4 and 5: Ensure all nodes (and in particular, the end node) are reachable from the start node
+    reachable: set[NodeID] = {StartNode.id}
+    to_visit = [StartNode.id]
+
+    while to_visit:
+        current_id = to_visit.pop()
+
+        # Add destinations from regular edges
+        for path in edges_by_source.get(current_id, []):
+            for item in path.items:
+                if isinstance(item, DestinationMarker):
+                    if item.destination_id not in reachable:
+                        reachable.add(item.destination_id)
+                        to_visit.append(item.destination_id)
+
+        # Add destinations from decision branches
+        current_node = nodes.get(current_id)
+        if isinstance(current_node, Decision):
+            for branch in current_node.branches:
+                for item in branch.path.items:
+                    if isinstance(item, DestinationMarker):
+                        if item.destination_id not in reachable:
+                            reachable.add(item.destination_id)
+                            to_visit.append(item.destination_id)
+
+    unreachable_nodes = [node_id for node_id in nodes if node_id not in reachable]
+    if unreachable_nodes:
+        raise GraphValidationError(
+            f'The following nodes are not reachable from the start node: {unreachable_nodes}.' + how_to_suppress
+        )
+
+
+def _flatten_paths(
+    nodes: dict[NodeID, AnyNode], edges: dict[NodeID, list[Path]]
+) -> tuple[dict[NodeID, AnyNode], dict[NodeID, list[Path]]]:
+    new_nodes = nodes.copy()
+    new_edges: dict[NodeID, list[Path]] = defaultdict(list)
+
+    paths_to_handle: list[tuple[NodeID, Path]] = []
+
+    def _split_at_first_fork(path: Path) -> tuple[Path, list[tuple[NodeID, Path]]]:
+        for i, item in enumerate(path.items):
+            if isinstance(item, MapMarker):
+                assert item.fork_id in nodes, 'This should have been added to the node during GraphBuilder.add'
+                upstream = Path(list(path.items[:i]) + [DestinationMarker(item.fork_id)])
+                downstream = Path(path.items[i + 1 :])
+                return upstream, [(item.fork_id, downstream)]
+
+            if isinstance(item, BroadcastMarker):
+                assert item.fork_id in nodes, 'This should have been added to the node during GraphBuilder.add'
+                upstream = Path(list(path.items[:i]) + [DestinationMarker(item.fork_id)])
+                return upstream, [(item.fork_id, p) for p in item.paths]
+        return path, []
+
+    for node in new_nodes.values():
+        if isinstance(node, Decision):
+            for branch in node.branches:
+                upstream, downstreams = _split_at_first_fork(branch.path)
+                branch.path = upstream
+                paths_to_handle.extend(downstreams)
+
+    for source_id, edges_from_source in edges.items():
+        for path in edges_from_source:
+            paths_to_handle.append((source_id, path))
+
+    while paths_to_handle:
+        source_id, path = paths_to_handle.pop()
+        upstream, downstreams = _split_at_first_fork(path)
+        new_edges[source_id].append(upstream)
+        paths_to_handle.extend(downstreams)
+
+    return new_nodes, dict(new_edges)
+
+
+def _normalize_forks(
+    nodes: dict[NodeID, AnyNode], edges: dict[NodeID, list[Path]]
+) -> tuple[dict[NodeID, AnyNode], dict[NodeID, list[Path]]]:
+    """Normalize the graph structure so only broadcast forks have multiple outgoing edges.
+
+    This function ensures that any node with multiple outgoing edges is converted
+    to use an explicit broadcast fork, simplifying the graph execution model.
+
+    Args:
+        nodes: The nodes in the graph
+        edges: The edges by source node
+
+    Returns:
+        A tuple of normalized nodes and edges
+    """
+    new_nodes = nodes.copy()
+    new_edges: dict[NodeID, list[Path]] = {}
+
+    paths_to_handle: list[Path] = []
+
+    for source_id, edges_from_source in edges.items():
+        paths_to_handle.extend(edges_from_source)
+
+        node = nodes[source_id]
+        if isinstance(node, Fork) and not node.is_map:
+            new_edges[source_id] = edges_from_source
+            continue  # broadcast fork; nothing to do
+        if len(edges_from_source) == 1:
+            new_edges[source_id] = edges_from_source
+            continue
+        new_fork = Fork[Any, Any](id=ForkID(NodeID(f'{node.id}_broadcast_fork')), is_map=False, downstream_join_id=None)
+        new_nodes[new_fork.id] = new_fork
+        new_edges[source_id] = [Path(items=[DestinationMarker(new_fork.id)])]
+        new_edges[new_fork.id] = edges_from_source
+
+    return new_nodes, new_edges
+
+
+def _collect_dominating_forks(
+    graph_nodes: dict[NodeID, AnyNode], graph_edges_by_source: dict[NodeID, list[Path]]
+) -> dict[JoinID, ParentFork[NodeID]]:
+    """Find the dominating fork for each join node in the graph.
+
+    This function analyzes the graph structure to find the parent fork that
+    dominates each join node, which is necessary for proper synchronization
+    during graph execution.
+
+    Args:
+        graph_nodes: All nodes in the graph
+        graph_edges_by_source: Edges organized by source node
+
+    Returns:
+        A mapping from join IDs to their parent fork information
+
+    Raises:
+        ValueError: If any join node lacks a dominating fork
+    """
+    nodes = set(graph_nodes)
+    start_ids: set[NodeID] = {StartNode.id}
+    edges: dict[NodeID, list[NodeID]] = defaultdict(list)
+
+    fork_ids: set[NodeID] = set(start_ids)
+    for source_id in nodes:
+        working_source_id = source_id
+        node = graph_nodes.get(source_id)
+
+        if isinstance(node, Fork):
+            fork_ids.add(node.id)
+
+        def _handle_path(path: Path, last_source_id: NodeID):
+            """Process a path and collect edges and fork information.
+
+            Args:
+                path: The path to process
+                last_source_id: The current source node ID
+            """
+            for item in path.items:  # pragma: no branch
+                # No need to handle MapMarker or BroadcastMarker here as these should have all been removed
+                # by the call to `_flatten_paths`
+                if isinstance(item, DestinationMarker):
+                    edges[last_source_id].append(item.destination_id)
+                    # Destinations should only ever occur as the last item in the list, so no need to update the working_source_id
+                    break
+
+        if isinstance(node, Decision):
+            for branch in node.branches:
+                _handle_path(branch.path, working_source_id)
+        else:
+            for path in graph_edges_by_source.get(source_id, []):
+                _handle_path(path, source_id)
+
+    finder = ParentForkFinder(
+        nodes=nodes,
+        start_ids=start_ids,
+        fork_ids=fork_ids,
+        edges=edges,
+    )
+
+    joins = [node for node in graph_nodes.values() if isinstance(node, Join)]
+    dominating_forks: dict[JoinID, ParentFork[NodeID]] = {}
+    for join in joins:
+        dominating_fork = finder.find_parent_fork(
+            join.id, parent_fork_id=join.parent_fork_id, prefer_closest=join.preferred_parent_fork == 'closest'
+        )
+        if dominating_fork is None:
+            rendered_mermaid_graph = build_mermaid_graph(graph_nodes, graph_edges_by_source).render()
+            raise GraphBuildingError(f"""A node in the graph is missing a dominating fork.
+
+For every Join J in the graph, there must be a Fork F between the StartNode and J satisfying:
+* Every path from the StartNode to J passes through F
+* There are no cycles in the graph including J that don't pass through F.
+In this case, F is called a "dominating fork" for J.
+
+This is used to determine when all tasks upstream of this Join are complete and we can proceed with execution.
+
+Mermaid diagram:
+{rendered_mermaid_graph}
+
+Join {join.id!r} in this graph has no dominating fork in this graph.""")
+        dominating_forks[join.id] = dominating_fork
+
+    return dominating_forks
+
+
+def _compute_intermediate_join_nodes(
+    nodes: dict[NodeID, AnyNode], parent_forks: dict[JoinID, ParentFork[NodeID]]
+) -> dict[JoinID, set[JoinID]]:
+    """Compute which joins have other joins as intermediate nodes.
+
+    A join J1 is an intermediate node of join J2 if J1 appears in J2's intermediate_nodes
+    (as computed relative to J2's parent fork).
+
+    This information is used to determine:
+    1. Which joins are "final" (have no other joins in their intermediate_nodes)
+    2. When selecting which reducer to proceed with when there are no active tasks
+
+    Args:
+        nodes: All nodes in the graph
+        parent_forks: Parent fork information for each join
+
+    Returns:
+        A mapping from each join to the set of joins that are intermediate to it
+    """
+    intermediate_join_nodes: dict[JoinID, set[JoinID]] = {}
+
+    for join_id, parent_fork in parent_forks.items():
+        intermediate_joins = set[JoinID]()
+        for intermediate_node_id in parent_fork.intermediate_nodes:
+            # Check if this intermediate node is also a join
+            intermediate_node = nodes.get(intermediate_node_id)
+            if isinstance(intermediate_node, Join):
+                # Add it regardless of whether it has the same parent fork
+                intermediate_joins.add(JoinID(intermediate_node_id))
+        intermediate_join_nodes[join_id] = intermediate_joins
+
+    return intermediate_join_nodes
+
+
+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 in nodes.keys():
+        replaced_node_id = replace_placeholder_id(node_id)
+        if replaced_node_id == node_id:
+            continue
+        counter[replaced_node_id] = count = counter[replaced_node_id] + 1
+        remapping[node_id] = NodeID(f'{replaced_node_id}_{count}' if count > 1 else replaced_node_id)
+    return remapping
+
+
+def _update_node_with_id_remapping(node: AnyNode, node_id_remapping: dict[NodeID, NodeID]) -> AnyNode:
+    # Note: it's a bit awkward that we mutate the provided nodes, but this is necessary to ensure that
+    # calls to `.as_node` reference the correct node_ids when relying on compatibility with the v1 API.
+    # We only mutate placeholder IDs so I _think_ this should generally be okay. I guess we can
+    # rework it more carefully if it causes issues in the future..
+    if isinstance(node, Step):
+        node.id = node_id_remapping.get(node.id, node.id)
+    elif isinstance(node, Join):
+        node.id = JoinID(node_id_remapping.get(node.id, node.id))
+    elif isinstance(node, Fork):
+        node.id = ForkID(node_id_remapping.get(node.id, node.id))
+        if node.downstream_join_id is not None:
+            node.downstream_join_id = JoinID(node_id_remapping.get(node.downstream_join_id, node.downstream_join_id))
+    elif isinstance(node, Decision):
+        node.id = node_id_remapping.get(node.id, node.id)
+        node.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:
+    # Note: we have already deepcopied the node provided 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.
+    for item in path.items:
+        if isinstance(item, MapMarker):
+            downstream_join_id = item.downstream_join_id
+            if downstream_join_id is not None:
+                item.downstream_join_id = JoinID(node_id_remapping.get(downstream_join_id, downstream_join_id))
+            item.fork_id = ForkID(node_id_remapping.get(item.fork_id, item.fork_id))
+        elif isinstance(item, BroadcastMarker):
+            item.fork_id = ForkID(node_id_remapping.get(item.fork_id, item.fork_id))
+            item.paths = [_update_path_with_id_remapping(p, node_id_remapping) for p in item.paths]
+        elif isinstance(item, DestinationMarker):
+            item.destination_id = node_id_remapping.get(item.destination_id, item.destination_id)
+    return path