pydantic_graph.beta.decision

Decision node implementation for conditional branching in graph execution.

This module provides the Decision node type and related classes for implementing conditional branching logic in parallel control flow graphs. Decision nodes allow the graph to choose different execution paths based on runtime conditions.

StateT module-attribute

StateT = TypeVar('StateT', infer_variance=True)

Type variable for graph state.

DepsT module-attribute

DepsT = TypeVar('DepsT', infer_variance=True)

Type variable for graph dependencies.

HandledT module-attribute

HandledT = TypeVar('HandledT', infer_variance=True)

Type variable used to track types handled by the branches of a Decision.

T module-attribute

T = TypeVar('T', infer_variance=True)

Generic type variable.

Decision dataclass

Bases: Generic[StateT, DepsT, HandledT]

Decision node for conditional branching in graph execution.

A Decision node evaluates conditions and routes execution to different branches based on the input data type or custom matching logic.

Source code in pydantic_graph/pydantic_graph/beta/decision.py

@dataclass(kw_only=True)
class Decision(Generic[StateT, DepsT, HandledT]):
    """Decision node for conditional branching in graph execution.

    A Decision node evaluates conditions and routes execution to different
    branches based on the input data type or custom matching logic.
    """

    id: NodeID
    """Unique identifier for this decision node."""

    branches: list[DecisionBranch[Any]]
    """List of branches that can be taken from this decision."""

    note: str | None
    """Optional documentation note for this decision."""

    def branch(self, branch: DecisionBranch[T]) -> Decision[StateT, DepsT, HandledT | T]:
        """Add a new branch to this decision.

        Args:
            branch: The branch to add to this decision.

        Returns:
            A new Decision with the additional branch.
        """
        return Decision(id=self.id, branches=self.branches + [branch], note=self.note)

    def _force_handled_contravariant(self, inputs: HandledT) -> Never:  # pragma: no cover
        """Forces this type to be contravariant in the HandledT type variable.

        This is an implementation detail of how we can type-check that all possible input types have
        been exhaustively covered.

        Args:
            inputs: Input data of handled types.

        Raises:
            RuntimeError: Always, as this method should never be executed.
        """
        raise RuntimeError('This method should never be called, it is just defined for typing purposes.')

id instance-attribute

id: NodeID

Unique identifier for this decision node.

branches instance-attribute

branches: list[DecisionBranch[Any]]

List of branches that can be taken from this decision.

note instance-attribute

note: str | None

Optional documentation note for this decision.

branch

branch(
    branch: DecisionBranch[T],
) -> Decision[StateT, DepsT, HandledT | T]

Add a new branch to this decision.

Parameters:

Name	Type	Description	Default
branch	DecisionBranch[T]	The branch to add to this decision.	required

Returns:

Type	Description
Decision[StateT, DepsT, HandledT | T]	A new Decision with the additional branch.

Source code in pydantic_graph/pydantic_graph/beta/decision.py

def branch(self, branch: DecisionBranch[T]) -> Decision[StateT, DepsT, HandledT | T]:
    """Add a new branch to this decision.

    Args:
        branch: The branch to add to this decision.

    Returns:
        A new Decision with the additional branch.
    """
    return Decision(id=self.id, branches=self.branches + [branch], note=self.note)

SourceT module-attribute

SourceT = TypeVar('SourceT', infer_variance=True)

Type variable for source data for a DecisionBranch.

DecisionBranch dataclass

Bases: Generic[SourceT]

Represents a single branch within a decision node.

Each branch defines the conditions under which it should be taken and the path to follow when those conditions are met.

Note: with the current design, it is actually critical that this class is invariant in SourceT for the sake of type-checking that inputs to a Decision are actually handled. See the # type: ignore comment in tests.graph.beta.test_graph_edge_cases.test_decision_no_matching_branch for an example of how this works.

Source code in pydantic_graph/pydantic_graph/beta/decision.py

@dataclass
class DecisionBranch(Generic[SourceT]):
    """Represents a single branch within a decision node.

    Each branch defines the conditions under which it should be taken
    and the path to follow when those conditions are met.

    Note: with the current design, it is actually _critical_ that this class is invariant in SourceT for the sake
    of type-checking that inputs to a Decision are actually handled. See the `# type: ignore` comment in
    `tests.graph.beta.test_graph_edge_cases.test_decision_no_matching_branch` for an example of how this works.
    """

    source: TypeOrTypeExpression[SourceT]
    """The expected type of data for this branch.

    This is necessary for exhaustiveness-checking when handling the inputs to a decision node."""

    matches: Callable[[Any], bool] | None
    """An optional predicate function used to determine whether input data matches this branch.

    If `None`, default logic is used which attempts to check the value for type-compatibility with the `source` type:
    * If `source` is `Any` or `object`, the branch will always match
    * If `source` is a `Literal` type, this branch will match if the value is one of the parametrizing literal values
    * If `source` is any other type, the value will be checked for matching using `isinstance`

    Inputs are tested against each branch of a decision node in order, and the path of the first matching branch is
    used to handle the input value.
    """

    path: Path
    """The execution path to follow when an input value matches this branch of a decision node.

    This can include transforming, mapping, and broadcasting the output before sending to the next node or nodes.

    The path can also include position-aware labels which are used when generating mermaid diagrams."""

    destinations: list[AnyDestinationNode]
    """The destination nodes that can be referenced by DestinationMarker in the path."""

source instance-attribute

source: TypeOrTypeExpression[SourceT]

The expected type of data for this branch.

This is necessary for exhaustiveness-checking when handling the inputs to a decision node.

matches instance-attribute

matches: Callable[[Any], bool] | None

An optional predicate function used to determine whether input data matches this branch.

If None, default logic is used which attempts to check the value for type-compatibility with the source type: * If source is Any or object, the branch will always match * If source is a Literal type, this branch will match if the value is one of the parametrizing literal values * If source is any other type, the value will be checked for matching using isinstance

Inputs are tested against each branch of a decision node in order, and the path of the first matching branch is used to handle the input value.

path instance-attribute

path: Path

The execution path to follow when an input value matches this branch of a decision node.

This can include transforming, mapping, and broadcasting the output before sending to the next node or nodes.

The path can also include position-aware labels which are used when generating mermaid diagrams.

destinations instance-attribute

destinations: list[AnyDestinationNode]

The destination nodes that can be referenced by DestinationMarker in the path.

OutputT module-attribute

OutputT = TypeVar('OutputT', infer_variance=True)

Type variable for the output data of a node.

NewOutputT module-attribute

NewOutputT = TypeVar('NewOutputT', infer_variance=True)

Type variable for transformed output.

DecisionBranchBuilder dataclass

Bases: 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, not created directly.

Source code in pydantic_graph/pydantic_graph/beta/decision.py

@dataclass(init=False)
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.
    """

    _decision: Decision[StateT, DepsT, HandledT]
    """The parent decision node."""
    _source: TypeOrTypeExpression[SourceT]
    """The expected source type for this branch."""
    _matches: Callable[[Any], bool] | None
    """Optional matching predicate."""

    _path_builder: PathBuilder[StateT, DepsT, OutputT]
    """Builder for the execution path."""

    def __init__(
        self,
        *,
        decision: Decision[StateT, DepsT, HandledT],
        source: TypeOrTypeExpression[SourceT],
        matches: Callable[[Any], bool] | None,
        path_builder: PathBuilder[StateT, DepsT, OutputT],
    ):
        # This manually-defined initializer is necessary due to https://github.com/python/mypy/issues/17623.
        self._decision = decision
        self._source = source
        self._matches = matches
        self._path_builder = path_builder

    def to(
        self,
        destination: DestinationNode[StateT, DepsT, OutputT] | type[BaseNode[StateT, DepsT, Any]],
        /,
        *extra_destinations: DestinationNode[StateT, DepsT, OutputT] | type[BaseNode[StateT, DepsT, Any]],
        fork_id: str | None = None,
    ) -> DecisionBranch[SourceT]:
        """Set the destination(s) for this branch.

        Args:
            destination: The primary destination node.
            *extra_destinations: Additional destination nodes.
            fork_id: Optional node ID to use for the resulting broadcast fork if multiple destinations are provided.

        Returns:
            A completed DecisionBranch with the specified destinations.
        """
        destination = get_origin(destination) or destination
        extra_destinations = tuple(get_origin(d) or d for d in extra_destinations)
        destinations = [(NodeStep(d) if inspect.isclass(d) else d) for d in (destination, *extra_destinations)]
        return DecisionBranch(
            source=self._source,
            matches=self._matches,
            path=self._path_builder.to(*destinations, fork_id=fork_id),
            destinations=destinations,
        )

    def broadcast(
        self, get_forks: Callable[[Self], Sequence[DecisionBranch[SourceT]]], /, *, fork_id: str | None = None
    ) -> DecisionBranch[SourceT]:
        """Broadcast this decision branch into multiple destinations.

        Args:
            get_forks: The callback that will return a sequence of decision branches to broadcast to.
            fork_id: Optional node ID to use for the resulting broadcast fork.

        Returns:
            A completed DecisionBranch with the specified destinations.
        """
        fork_decision_branches = get_forks(self)
        new_paths = [b.path for b in fork_decision_branches]
        if not new_paths:
            raise GraphBuildingError(f'The call to {get_forks} returned no branches, but must return at least one.')
        path = self._path_builder.broadcast(new_paths, fork_id=fork_id)
        destinations = [d for fdp in fork_decision_branches for d in fdp.destinations]
        return DecisionBranch(source=self._source, matches=self._matches, path=path, destinations=destinations)

    def transform(
        self, func: TransformFunction[StateT, DepsT, OutputT, NewOutputT], /
    ) -> DecisionBranchBuilder[StateT, DepsT, NewOutputT, SourceT, HandledT]:
        """Apply a transformation to the branch's output.

        Args:
            func: Transformation function to apply.

        Returns:
            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),
        )

    def map(
        self: DecisionBranchBuilder[StateT, DepsT, Iterable[T], SourceT, HandledT]
        | DecisionBranchBuilder[StateT, DepsT, AsyncIterable[T], SourceT, HandledT],
        *,
        fork_id: str | None = None,
        downstream_join_id: str | None = None,
    ) -> DecisionBranchBuilder[StateT, DepsT, T, SourceT, HandledT]:
        """Spread the branch's output.

        To do this, the current output must be iterable, and any subsequent steps in the path being built for this
        branch will be applied to each item of the current output in parallel.

        Args:
            fork_id: Optional ID for the fork, defaults to a generated value
            downstream_join_id: Optional ID of a downstream join node which is involved when mapping empty iterables

        Returns:
            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),
        )

    def label(self, label: str) -> DecisionBranchBuilder[StateT, DepsT, OutputT, SourceT, HandledT]:
        """Apply a label to the branch at the current point in the path being built.

        These labels are only used in generated mermaid diagrams.

        Args:
            label: The label to apply.

        Returns:
            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),
        )

to

to(
    destination: (
        DestinationNode[StateT, DepsT, OutputT]
        | type[BaseNode[StateT, DepsT, Any]]
    ),
    /,
    *extra_destinations: DestinationNode[
        StateT, DepsT, OutputT
    ]
    | type[BaseNode[StateT, DepsT, Any]],
    fork_id: str | None = None,
) -> DecisionBranch[SourceT]

Set the destination(s) for this branch.

Parameters:

Name	Type	Description	Default
destination	DestinationNode[StateT, DepsT, OutputT] | type[BaseNode[StateT, DepsT, Any]]	The primary destination node.	required
*extra_destinations	DestinationNode[StateT, DepsT, OutputT] | type[BaseNode[StateT, DepsT, Any]]	Additional destination nodes.	()
fork_id	str | None	Optional node ID to use for the resulting broadcast fork if multiple destinations are provided.	None

Returns:

Type	Description
DecisionBranch[SourceT]	A completed DecisionBranch with the specified destinations.

Source code in pydantic_graph/pydantic_graph/beta/decision.py

def to(
    self,
    destination: DestinationNode[StateT, DepsT, OutputT] | type[BaseNode[StateT, DepsT, Any]],
    /,
    *extra_destinations: DestinationNode[StateT, DepsT, OutputT] | type[BaseNode[StateT, DepsT, Any]],
    fork_id: str | None = None,
) -> DecisionBranch[SourceT]:
    """Set the destination(s) for this branch.

    Args:
        destination: The primary destination node.
        *extra_destinations: Additional destination nodes.
        fork_id: Optional node ID to use for the resulting broadcast fork if multiple destinations are provided.

    Returns:
        A completed DecisionBranch with the specified destinations.
    """
    destination = get_origin(destination) or destination
    extra_destinations = tuple(get_origin(d) or d for d in extra_destinations)
    destinations = [(NodeStep(d) if inspect.isclass(d) else d) for d in (destination, *extra_destinations)]
    return DecisionBranch(
        source=self._source,
        matches=self._matches,
        path=self._path_builder.to(*destinations, fork_id=fork_id),
        destinations=destinations,
    )

broadcast

broadcast(
    get_forks: Callable[
        [Self], Sequence[DecisionBranch[SourceT]]
    ],
    /,
    *,
    fork_id: str | None = None,
) -> DecisionBranch[SourceT]

Broadcast this decision branch into multiple destinations.

Parameters:

Name	Type	Description	Default
get_forks	Callable[[Self], Sequence[DecisionBranch[SourceT]]]	The callback that will return a sequence of decision branches to broadcast to.	required
fork_id	str | None	Optional node ID to use for the resulting broadcast fork.	None

Returns:

Type	Description
DecisionBranch[SourceT]	A completed DecisionBranch with the specified destinations.

Source code in pydantic_graph/pydantic_graph/beta/decision.py

def broadcast(
    self, get_forks: Callable[[Self], Sequence[DecisionBranch[SourceT]]], /, *, fork_id: str | None = None
) -> DecisionBranch[SourceT]:
    """Broadcast this decision branch into multiple destinations.

    Args:
        get_forks: The callback that will return a sequence of decision branches to broadcast to.
        fork_id: Optional node ID to use for the resulting broadcast fork.

    Returns:
        A completed DecisionBranch with the specified destinations.
    """
    fork_decision_branches = get_forks(self)
    new_paths = [b.path for b in fork_decision_branches]
    if not new_paths:
        raise GraphBuildingError(f'The call to {get_forks} returned no branches, but must return at least one.')
    path = self._path_builder.broadcast(new_paths, fork_id=fork_id)
    destinations = [d for fdp in fork_decision_branches for d in fdp.destinations]
    return DecisionBranch(source=self._source, matches=self._matches, path=path, destinations=destinations)

transform

transform(
    func: TransformFunction[
        StateT, DepsT, OutputT, NewOutputT
    ],
) -> DecisionBranchBuilder[
    StateT, DepsT, NewOutputT, SourceT, HandledT
]

Apply a transformation to the branch's output.

Parameters:

Name	Type	Description	Default
func	TransformFunction[StateT, DepsT, OutputT, NewOutputT]	Transformation function to apply.	required

Returns:

Type	Description
DecisionBranchBuilder[StateT, DepsT, NewOutputT, SourceT, HandledT]	A new DecisionBranchBuilder where the provided transform is applied prior to generating the final output.

Source code in pydantic_graph/pydantic_graph/beta/decision.py

def transform(
    self, func: TransformFunction[StateT, DepsT, OutputT, NewOutputT], /
) -> DecisionBranchBuilder[StateT, DepsT, NewOutputT, SourceT, HandledT]:
    """Apply a transformation to the branch's output.

    Args:
        func: Transformation function to apply.

    Returns:
        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),
    )

map

map(
    *,
    fork_id: str | None = None,
    downstream_join_id: str | None = None
) -> DecisionBranchBuilder[
    StateT, DepsT, T, SourceT, HandledT
]

Spread the branch's output.

To do this, the current output must be iterable, and any subsequent steps in the path being built for this branch will be applied to each item of the current output in parallel.

Parameters:

Name	Type	Description	Default
fork_id	str | None	Optional ID for the fork, defaults to a generated value	None
downstream_join_id	str | None	Optional ID of a downstream join node which is involved when mapping empty iterables	None

Returns:

Type	Description
DecisionBranchBuilder[StateT, DepsT, T, SourceT, HandledT]	A new DecisionBranchBuilder where mapping is performed prior to generating the final output.

Source code in pydantic_graph/pydantic_graph/beta/decision.py

def map(
    self: DecisionBranchBuilder[StateT, DepsT, Iterable[T], SourceT, HandledT]
    | DecisionBranchBuilder[StateT, DepsT, AsyncIterable[T], SourceT, HandledT],
    *,
    fork_id: str | None = None,
    downstream_join_id: str | None = None,
) -> DecisionBranchBuilder[StateT, DepsT, T, SourceT, HandledT]:
    """Spread the branch's output.

    To do this, the current output must be iterable, and any subsequent steps in the path being built for this
    branch will be applied to each item of the current output in parallel.

    Args:
        fork_id: Optional ID for the fork, defaults to a generated value
        downstream_join_id: Optional ID of a downstream join node which is involved when mapping empty iterables

    Returns:
        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),
    )

label

label(
    label: str,
) -> DecisionBranchBuilder[
    StateT, DepsT, OutputT, SourceT, HandledT
]

Apply a label to the branch at the current point in the path being built.

These labels are only used in generated mermaid diagrams.

Parameters:

Name	Type	Description	Default
label	str	The label to apply.	required

Returns:

Type	Description
DecisionBranchBuilder[StateT, DepsT, OutputT, SourceT, HandledT]	A new DecisionBranchBuilder where the label has been applied at the end of the current path being built.

Source code in pydantic_graph/pydantic_graph/beta/decision.py

def label(self, label: str) -> DecisionBranchBuilder[StateT, DepsT, OutputT, SourceT, HandledT]:
    """Apply a label to the branch at the current point in the path being built.

    These labels are only used in generated mermaid diagrams.

    Args:
        label: The label to apply.

    Returns:
        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),
    )