pydantic-graph 1.3.0__py3-none-any.whl → 1.12.0__py3-none-any.whl

pydantic_graph/beta/paths.py

@@ -0,0 +1,421 @@
+"""Path and edge definition for graph navigation.
+
+This module provides the building blocks for defining paths through a graph,
+including transformations, maps, broadcasts, and routing to destinations.
+Paths enable complex data flow patterns in graph execution.
+"""
+
+from __future__ import annotations
+
+import inspect
+from collections.abc import AsyncIterable, Callable, Iterable, Sequence
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Generic, get_origin
+
+from typing_extensions import Protocol, Self, TypeAliasType, TypeVar
+
+from pydantic_graph import BaseNode
+from pydantic_graph.beta.id_types import ForkID, JoinID, NodeID, generate_placeholder_node_id
+from pydantic_graph.beta.step import NodeStep, StepContext
+from pydantic_graph.exceptions import GraphBuildingError
+
+StateT = TypeVar('StateT', infer_variance=True)
+DepsT = TypeVar('DepsT', infer_variance=True)
+OutputT = TypeVar('OutputT', infer_variance=True)
+InputT = TypeVar('InputT', infer_variance=True)
+T = TypeVar('T')
+
+if TYPE_CHECKING:
+    from pydantic_graph.beta.node_types import AnyDestinationNode, DestinationNode, SourceNode
+
+
+class TransformFunction(Protocol[StateT, DepsT, InputT, OutputT]):
+    """Protocol for step functions that can be executed in the graph.
+
+    Transform functions are sync callables that receive a step context and return
+    a result. This protocol enables serialization and deserialization of step
+    calls similar to how evaluators work.
+
+    This is very similar to a StepFunction, but must be sync instead of async.
+
+    Type Parameters:
+        StateT: The type of the graph state
+        DepsT: The type of the dependencies
+        InputT: The type of the input data
+        OutputT: The type of the output data
+    """
+
+    def __call__(self, ctx: StepContext[StateT, DepsT, InputT]) -> OutputT:
+        """Execute the step function with the given context.
+
+        Args:
+            ctx: The step context containing state, dependencies, and inputs
+
+        Returns:
+            An awaitable that resolves to the step's output
+        """
+        raise NotImplementedError
+
+
+@dataclass
+class TransformMarker:
+    """A marker indicating a data transformation step in a path.
+
+    Transform markers wrap step functions that modify data as it flows
+    through the graph path.
+    """
+
+    transform: TransformFunction[Any, Any, Any, Any]
+    """The step function that performs the transformation."""
+
+
+@dataclass
+class MapMarker:
+    """A marker indicating that iterable data should be map across parallel paths.
+
+    Spread markers take iterable input and create parallel execution paths
+    for each item in the iterable.
+    """
+
+    fork_id: ForkID
+    """Unique identifier for the fork created by this map operation."""
+    downstream_join_id: JoinID | None
+    """Optional identifier of a downstream join node that should be jumped to if mapping an empty iterable."""
+
+
+@dataclass
+class BroadcastMarker:
+    """A marker indicating that data should be broadcast to multiple parallel paths.
+
+    Broadcast markers create multiple parallel execution paths, sending the
+    same input data to each path.
+    """
+
+    paths: Sequence[Path]
+    """The parallel paths that will receive the broadcast data."""
+
+    fork_id: ForkID
+    """Unique identifier for the fork created by this broadcast operation."""
+
+
+@dataclass
+class LabelMarker:
+    """A marker providing a human-readable label for a path segment.
+
+    Label markers are used for debugging, visualization, and documentation
+    purposes to provide meaningful names for path segments.
+    """
+
+    label: str
+    """The human-readable label for this path segment."""
+
+
+@dataclass
+class DestinationMarker:
+    """A marker indicating the target destination node for a path.
+
+    Destination markers specify where data should be routed at the end
+    of a path execution.
+    """
+
+    destination_id: NodeID
+    """The unique identifier of the destination node."""
+
+
+PathItem = TypeAliasType('PathItem', TransformMarker | MapMarker | BroadcastMarker | LabelMarker | DestinationMarker)
+"""Type alias for any item that can appear in a path sequence."""
+
+
+@dataclass
+class Path:
+    """A sequence of path items defining data flow through the graph.
+
+    Paths represent the route that data takes through the graph, including
+    transformations, forks, and routing decisions.
+    """
+
+    items: list[PathItem]
+    """The sequence of path items that define this path."""
+
+    @property
+    def last_fork(self) -> BroadcastMarker | MapMarker | None:
+        """Get the most recent fork or map marker in this path.
+
+        Returns:
+            The last BroadcastMarker or MapMarker in the path, or None if no forks exist
+        """
+        for item in reversed(self.items):
+            if isinstance(item, BroadcastMarker | MapMarker):
+                return item
+        return None
+
+    @property
+    def next_path(self) -> Path:
+        """Create a new path with the first item removed.
+
+        Returns:
+            A new Path with all items except the first one
+        """
+        return Path(self.items[1:])
+
+
+@dataclass
+class PathBuilder(Generic[StateT, DepsT, OutputT]):
+    """A builder for constructing paths with method chaining.
+
+    PathBuilder provides a fluent interface for creating paths by chaining
+    operations like transforms, maps, and routing to destinations.
+
+    Type Parameters:
+        StateT: The type of the graph state
+        DepsT: The type of the dependencies
+        OutputT: The type of the current data in the path
+    """
+
+    working_items: Sequence[PathItem]
+    """The accumulated sequence of path items being built."""
+
+    def to(
+        self,
+        destination: DestinationNode[StateT, DepsT, OutputT],
+        /,
+        *extra_destinations: DestinationNode[StateT, DepsT, OutputT],
+        fork_id: str | None = None,
+    ) -> Path:
+        """Route the path to one or more destination nodes.
+
+        Args:
+            destination: The primary destination node
+            *extra_destinations: Additional destination nodes (creates a broadcast)
+            fork_id: Optional ID for the fork created when multiple destinations are specified
+
+        Returns:
+            A complete Path ending at the specified destination(s)
+        """
+        if extra_destinations:
+            next_item = BroadcastMarker(
+                paths=[Path(items=[DestinationMarker(d.id)]) for d in (destination,) + extra_destinations],
+                fork_id=ForkID(NodeID(fork_id or generate_placeholder_node_id('broadcast'))),
+            )
+        else:
+            next_item = DestinationMarker(destination.id)
+        return Path(items=[*self.working_items, next_item])
+
+    def broadcast(self, forks: Sequence[Path], /, *, fork_id: str | None = None) -> Path:
+        """Create a fork that broadcasts data to multiple parallel paths.
+
+        Args:
+            forks: The sequence of paths to run in parallel
+            fork_id: Optional ID for the fork, defaults to a generated value
+
+        Returns:
+            A complete Path that forks to the specified parallel paths
+        """
+        next_item = BroadcastMarker(
+            paths=forks, fork_id=ForkID(NodeID(fork_id or generate_placeholder_node_id('broadcast')))
+        )
+        return Path(items=[*self.working_items, next_item])
+
+    def transform(self, func: TransformFunction[StateT, DepsT, OutputT, T], /) -> PathBuilder[StateT, DepsT, T]:
+        """Add a transformation step to the path.
+
+        Args:
+            func: The step function that will transform the data
+
+        Returns:
+            A new PathBuilder with the transformation added
+        """
+        next_item = TransformMarker(func)
+        return PathBuilder[StateT, DepsT, T](working_items=[*self.working_items, next_item])
+
+    def map(
+        self: PathBuilder[StateT, DepsT, Iterable[T]] | PathBuilder[StateT, DepsT, AsyncIterable[T]],
+        *,
+        fork_id: str | None = None,
+        downstream_join_id: str | None = None,
+    ) -> PathBuilder[StateT, DepsT, T]:
+        """Spread iterable data across parallel execution paths.
+
+        This method can only be called when the current output type is iterable.
+        It creates parallel paths for each item in the iterable.
+
+        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 PathBuilder that operates on individual items from the iterable
+        """
+        next_item = MapMarker(
+            fork_id=ForkID(NodeID(fork_id or generate_placeholder_node_id('map'))),
+            downstream_join_id=JoinID(downstream_join_id) if downstream_join_id is not None else None,
+        )
+        return PathBuilder[StateT, DepsT, T](working_items=[*self.working_items, next_item])
+
+    def label(self, label: str, /) -> PathBuilder[StateT, DepsT, OutputT]:
+        """Add a human-readable label to this point in the path.
+
+        Args:
+            label: The label to add for documentation/debugging purposes
+
+        Returns:
+            A new PathBuilder with the label added
+        """
+        next_item = LabelMarker(label)
+        return PathBuilder[StateT, DepsT, OutputT](working_items=[*self.working_items, next_item])
+
+
+@dataclass(init=False)
+class EdgePath(Generic[StateT, DepsT]):
+    """A complete edge connecting source nodes to destinations via a path.
+
+    EdgePath represents a complete connection in the graph, specifying the
+    source nodes, the path that data follows, and the destination nodes.
+    """
+
+    _sources: Sequence[SourceNode[StateT, DepsT, Any]]
+    """The source nodes that provide data to this edge."""
+    path: Path
+    """The path that data follows through the graph."""
+    destinations: list[AnyDestinationNode]
+    """The destination nodes that can be referenced by DestinationMarker in the path."""
+
+    def __init__(
+        self, sources: Sequence[SourceNode[StateT, DepsT, Any]], path: Path, destinations: list[AnyDestinationNode]
+    ):
+        self._sources = sources
+        self.path = path
+        self.destinations = destinations
+
+    @property
+    def sources(self) -> Sequence[SourceNode[StateT, DepsT, Any]]:
+        return self._sources
+
+
+class EdgePathBuilder(Generic[StateT, DepsT, OutputT]):
+    """A builder for constructing complete edge paths with method chaining.
+
+    EdgePathBuilder combines source nodes with path building capabilities
+    to create complete edge definitions. It cannot use dataclass due to
+    type variance issues.
+
+    Type Parameters:
+        StateT: The type of the graph state
+        DepsT: The type of the dependencies
+        OutputT: The type of the current data in the path
+    """
+
+    def __init__(
+        self, sources: Sequence[SourceNode[StateT, DepsT, Any]], path_builder: PathBuilder[StateT, DepsT, OutputT]
+    ):
+        """Initialize an edge path builder.
+
+        Args:
+            sources: The source nodes for this edge path
+            path_builder: The path builder for defining the data flow
+        """
+        self.sources = sources
+        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,
+    ) -> EdgePath[StateT, DepsT]:
+        """Complete the edge path by routing to destination nodes.
+
+        Args:
+            destination: Either a destination node or a function that generates edge paths
+            *extra_destinations: Additional destination nodes (creates a broadcast)
+            fork_id: Optional ID for the fork created when multiple destinations are specified
+
+        Returns:
+            A complete EdgePath connecting sources to destinations
+        """
+        # `type[BaseNode[StateT, DepsT, Any]]` could actually be a `typing._GenericAlias` like `pydantic_ai._agent_graph.UserPromptNode[~DepsT, ~OutputT]`,
+        # so we get the origin to get to the actual class
+        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 EdgePath(
+            sources=self.sources,
+            path=self._path_builder.to(destinations[0], *destinations[1:], fork_id=fork_id),
+            destinations=destinations,
+        )
+
+    def broadcast(
+        self, get_forks: Callable[[Self], Sequence[EdgePath[StateT, DepsT]]], /, *, fork_id: str | None = None
+    ) -> EdgePath[StateT, DepsT]:
+        """Broadcast this EdgePathBuilder into multiple destinations.
+
+        Args:
+            get_forks: The callback that will return a sequence of EdgePaths to broadcast to.
+            fork_id: Optional node ID to use for the resulting broadcast fork.
+
+        Returns:
+            A completed EdgePath with the specified destinations.
+        """
+        new_edge_paths = get_forks(self)
+        new_paths = [Path(x.path.items) for x in new_edge_paths]
+        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 ep in new_edge_paths for d in ep.destinations]
+        return EdgePath(
+            sources=self.sources,
+            path=path,
+            destinations=destinations,
+        )
+
+    def map(
+        self: EdgePathBuilder[StateT, DepsT, Iterable[T]] | EdgePathBuilder[StateT, DepsT, AsyncIterable[T]],
+        *,
+        fork_id: str | None = None,
+        downstream_join_id: JoinID | None = None,
+    ) -> EdgePathBuilder[StateT, DepsT, T]:
+        """Spread iterable data across parallel execution paths.
+
+        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 EdgePathBuilder that operates on individual items from the iterable
+        """
+        if len(self.sources) > 1:
+            # The current implementation mishandles this because you get one copy of each edge
+            # from the MapMarker to its destination for each source, resulting in unintentional multiple execution.
+            # I suspect this is fixable without a major refactor, though it's not clear to me what the ideal behavior
+            # would be. But for now, it's definitely easiest to just raise an error for this.
+            raise NotImplementedError(
+                'Map is not currently supported with multiple source nodes.'
+                ' You can work around this by just creating a separate edge for each source.'
+            )
+        return EdgePathBuilder(
+            sources=self.sources,
+            path_builder=self._path_builder.map(fork_id=fork_id, downstream_join_id=downstream_join_id),
+        )
+
+    def transform(self, func: TransformFunction[StateT, DepsT, OutputT, T], /) -> EdgePathBuilder[StateT, DepsT, T]:
+        """Add a transformation step to the edge path.
+
+        Args:
+            func: The step function that will transform the data
+
+        Returns:
+            A new EdgePathBuilder with the transformation added
+        """
+        return EdgePathBuilder(sources=self.sources, path_builder=self._path_builder.transform(func))
+
+    def label(self, label: str) -> EdgePathBuilder[StateT, DepsT, OutputT]:
+        """Add a human-readable label to this point in the edge path.
+
+        Args:
+            label: The label to add for documentation/debugging purposes
+
+        Returns:
+            A new EdgePathBuilder with the label added
+        """
+        return EdgePathBuilder(sources=self.sources, path_builder=self._path_builder.label(label))

pydantic_graph/beta/step.py

@@ -0,0 +1,253 @@
+"""Step-based graph execution components.
+
+This module provides the core abstractions for step-based graph execution,
+including step contexts, step functions, and step nodes that bridge between
+the v1 and v2 graph execution systems.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Awaitable
+from dataclasses import dataclass
+from typing import Any, Generic, Protocol, cast, get_origin, overload
+
+from typing_extensions import TypeVar
+
+from pydantic_graph.beta.id_types import NodeID
+from pydantic_graph.nodes import BaseNode, End, GraphRunContext
+
+StateT = TypeVar('StateT', infer_variance=True)
+DepsT = TypeVar('DepsT', infer_variance=True)
+InputT = TypeVar('InputT', infer_variance=True)
+OutputT = TypeVar('OutputT', infer_variance=True)
+
+
+@dataclass(init=False)
+class StepContext(Generic[StateT, DepsT, InputT]):
+    """Context information passed to step functions during graph execution.
+
+    The step context provides access to the current graph state, dependencies, and input data for a step.
+
+    Type Parameters:
+        StateT: The type of the graph state
+        DepsT: The type of the dependencies
+        InputT: The type of the input data
+    """
+
+    _state: StateT
+    """The current graph state."""
+    _deps: DepsT
+    """The graph run dependencies."""
+    _inputs: InputT
+    """The input data for this step."""
+
+    def __init__(self, *, state: StateT, deps: DepsT, inputs: InputT):
+        self._state = state
+        self._deps = deps
+        self._inputs = inputs
+
+    @property
+    def state(self) -> StateT:
+        return self._state
+
+    @property
+    def deps(self) -> DepsT:
+        return self._deps
+
+    @property
+    def inputs(self) -> InputT:
+        """The input data for this step.
+
+        This must be a property to ensure correct variance behavior
+        """
+        return self._inputs
+
+
+class StepFunction(Protocol[StateT, DepsT, InputT, OutputT]):
+    """Protocol for step functions that can be executed in the graph.
+
+    Step functions are async callables that receive a step context and return a result.
+
+    Type Parameters:
+        StateT: The type of the graph state
+        DepsT: The type of the dependencies
+        InputT: The type of the input data
+        OutputT: The type of the output data
+    """
+
+    def __call__(self, ctx: StepContext[StateT, DepsT, InputT]) -> Awaitable[OutputT]:
+        """Execute the step function with the given context.
+
+        Args:
+            ctx: The step context containing state, dependencies, and inputs
+
+        Returns:
+            An awaitable that resolves to the step's output
+        """
+        raise NotImplementedError
+
+
+class StreamFunction(Protocol[StateT, DepsT, InputT, OutputT]):
+    """Protocol for stream functions that can be executed in the graph.
+
+    Stream functions are async callables that receive a step context and return an async iterator.
+
+    Type Parameters:
+        StateT: The type of the graph state
+        DepsT: The type of the dependencies
+        InputT: The type of the input data
+        OutputT: The type of the output data
+    """
+
+    def __call__(self, ctx: StepContext[StateT, DepsT, InputT]) -> AsyncIterator[OutputT]:
+        """Execute the stream function with the given context.
+
+        Args:
+            ctx: The step context containing state, dependencies, and inputs
+
+        Returns:
+            An async iterator yielding the streamed output
+        """
+        raise NotImplementedError
+        yield
+
+
+AnyStepFunction = StepFunction[Any, Any, Any, Any]
+"""Type alias for a step function with any type parameters."""
+
+
+@dataclass(init=False)
+class Step(Generic[StateT, DepsT, InputT, OutputT]):
+    """A step in the graph execution that wraps a step function.
+
+    Steps represent individual units of execution in the graph, encapsulating
+    a step function along with metadata like ID and label.
+
+    Type Parameters:
+        StateT: The type of the graph state
+        DepsT: The type of the dependencies
+        InputT: The type of the input data
+        OutputT: The type of the output data
+    """
+
+    id: NodeID
+    """Unique identifier for this step."""
+    _call: StepFunction[StateT, DepsT, InputT, OutputT]
+    """The step function to execute."""
+    label: str | None
+    """Optional human-readable label for this step."""
+
+    def __init__(self, *, id: NodeID, call: StepFunction[StateT, DepsT, InputT, OutputT], label: str | None = None):
+        self.id = id
+        self._call = call
+        self.label = label
+
+    @property
+    def call(self) -> StepFunction[StateT, DepsT, InputT, OutputT]:
+        """The step function to execute. This needs to be a property for proper variance inference."""
+        return self._call
+
+    @overload
+    def as_node(self, inputs: None = None) -> StepNode[StateT, DepsT]: ...
+
+    @overload
+    def as_node(self, inputs: InputT) -> StepNode[StateT, DepsT]: ...
+
+    def as_node(self, inputs: InputT | None = None) -> StepNode[StateT, DepsT]:
+        """Create a step node with bound inputs.
+
+        Args:
+            inputs: The input data to bind to this step, or None
+
+        Returns:
+            A [`StepNode`][pydantic_graph.beta.step.StepNode] with this step and the bound inputs
+        """
+        return StepNode(self, inputs)
+
+
+@dataclass
+class StepNode(BaseNode[StateT, DepsT, Any]):
+    """A base node that represents a step with bound inputs.
+
+    StepNode bridges between the v1 and v2 graph execution systems by wrapping
+    a [`Step`][pydantic_graph.beta.step.Step] with bound inputs in a BaseNode interface.
+    It is not meant to be run directly but rather used to indicate transitions
+    to v2-style steps.
+    """
+
+    step: Step[StateT, DepsT, Any, Any]
+    """The step to execute."""
+
+    inputs: Any
+    """The inputs bound to this step."""
+
+    async def run(self, ctx: GraphRunContext[StateT, DepsT]) -> BaseNode[StateT, DepsT, Any] | End[Any]:
+        """Attempt to run the step node.
+
+        Args:
+            ctx: The graph execution context
+
+        Returns:
+            The result of step execution
+
+        Raises:
+            NotImplementedError: Always raised as StepNode is not meant to be run directly
+        """
+        raise NotImplementedError(
+            '`StepNode` is not meant to be run directly, it is meant to be used in `BaseNode` subclasses to indicate a transition to v2-style steps.'
+        )
+
+
+# Note: we should make this into a frozen dataclass if https://github.com/python/mypy/issues/17623 gets resolved
+# Right now, it cannot be because that breaks variance inference in Python 3.13 due to __replace__
+class NodeStep(Step[StateT, DepsT, Any, BaseNode[StateT, DepsT, Any] | End[Any]]):
+    """A step that wraps a BaseNode type for execution.
+
+    NodeStep allows v1-style BaseNode classes to be used as steps in the
+    v2 graph execution system. It validates that the input is of the expected
+    node type and runs it with the appropriate graph context.
+    """
+
+    node_type: type[BaseNode[StateT, DepsT, Any]]
+    """The BaseNode type this step executes."""
+
+    def __init__(
+        self,
+        node_type: type[BaseNode[StateT, DepsT, Any]],
+        *,
+        id: NodeID | None = None,
+        label: str | None = None,
+    ):
+        """Initialize a node step.
+
+        Args:
+            node_type: The BaseNode class this step will execute
+            id: Optional unique identifier, defaults to the node's get_node_id()
+            label: Optional human-readable label for this step
+        """
+        super().__init__(
+            id=id or NodeID(node_type.get_node_id()),
+            call=self._call_node,
+            label=label,
+        )
+        # `type[BaseNode[StateT, DepsT, Any]]` could actually be a `typing._GenericAlias` like `pydantic_ai._agent_graph.UserPromptNode[~DepsT, ~OutputT]`,
+        # so we get the origin to get to the actual class
+        self.node_type = get_origin(node_type) or node_type
+
+    async def _call_node(self, ctx: StepContext[StateT, DepsT, Any]) -> BaseNode[StateT, DepsT, Any] | End[Any]:
+        """Execute the wrapped node with the step context.
+
+        Args:
+            ctx: The step context containing the node instance to run
+
+        Returns:
+            The result of running the node, either another BaseNode or End
+
+        Raises:
+            ValueError: If the input node is not of the expected type
+        """
+        node = ctx.inputs
+        if not isinstance(node, self.node_type):
+            raise ValueError(f'Node {node} is not of type {self.node_type}')  # pragma: no cover
+        node = cast(BaseNode[StateT, DepsT, Any], node)
+        return await node.run(GraphRunContext(state=ctx.state, deps=ctx.deps))