Skip to content

pydantic_graph.beta.step

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.

StepContext

Bases: 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.

Attributes

inputs

The input data for this step.

This must be a property to ensure correct variance behavior

Type: InputT

StepFunction

Bases: 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.

Methods

__call__
def __call__(ctx: StepContext[StateT, DepsT, InputT]) -> Awaitable[OutputT]

Execute the step function with the given context.

Returns

Awaitable[OutputT] — An awaitable that resolves to the step’s output

Parameters

ctx : StepContext[StateT, DepsT, InputT]

The step context containing state, dependencies, and inputs

StreamFunction

Bases: 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.

Methods

__call__
def __call__(ctx: StepContext[StateT, DepsT, InputT]) -> AsyncIterator[OutputT]

Execute the stream function with the given context.

Returns

AsyncIterator[OutputT] — An async iterator yielding the streamed output

Parameters

ctx : StepContext[StateT, DepsT, InputT]

The step context containing state, dependencies, and inputs

Step

Bases: 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.

Attributes

id

Unique identifier for this step.

Type: NodeID Default: id

label

Optional human-readable label for this step.

Type: str | None Default: label

call

The step function to execute. This needs to be a property for proper variance inference.

Type: StepFunction[StateT, DepsT, InputT, OutputT]

Methods

as_node
def as_node(inputs: None = None) -> StepNode[StateT, DepsT]
def as_node(inputs: InputT) -> StepNode[StateT, DepsT]

Create a step node with bound inputs.

Returns

StepNode[StateT, DepsT] — A StepNode with this step and the bound inputs

Parameters

inputs : InputT | None Default: None

The input data to bind to this step, or None

StepNode

Bases: 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 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.

Attributes

step

The step to execute.

Type: Step[StateT, DepsT, Any, Any]

inputs

The inputs bound to this step.

Type: Any

Methods

run

@async

def run(ctx: GraphRunContext[StateT, DepsT]) -> BaseNode[StateT, DepsT, Any] | End[Any]

Attempt to run the step node.

Returns

BaseNode[StateT, DepsT, Any] | End[Any] — The result of step execution

Parameters

ctx : GraphRunContext[StateT, DepsT]

The graph execution context

Raises
  • NotImplementedError — Always raised as StepNode is not meant to be run directly

NodeStep

Bases: 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.

Attributes

node_type

The BaseNode type this step executes.

Type: type[BaseNode[StateT, DepsT, Any]] Default: get_origin(node_type) or node_type

Methods

__init__
def __init__(
    node_type: type[BaseNode[StateT, DepsT, Any]],
    id: NodeID | None = None,
    label: str | None = None,
)

Initialize a node step.

Parameters

node_type : type[BaseNode[StateT, DepsT, Any]]

The BaseNode class this step will execute

id : NodeID | None Default: None

Optional unique identifier, defaults to the node’s get_node_id()

label : str | None Default: None

Optional human-readable label for this step

AnyStepFunction

Type alias for a step function with any type parameters.

Default: StepFunction[Any, Any, Any, Any]