AxFlowExecutionPlanner

Defined in: https://github.com/ax-llm/ax/blob/9a5a7060a48f9eef46efc680b0cdf6b42bff5df2/src/ax/flow/executionPlanner.ts#L27

Builds and manages the execution plan with automatic parallelization.

This class is the core of AxFlow’s performance optimization system. It analyzes the dependency relationships between steps and creates an optimized execution plan that maximizes parallelism while ensuring correct execution order.

Key responsibilities:

1. Dependency Analysis: Tracks what fields each step depends on and produces
2. Parallel Grouping: Groups independent steps that can run simultaneously
3. Execution Optimization: Creates optimized execution functions that run parallel groups concurrently
4. Signature Inference: Provides data for automatic signature generation

The planner works by building a directed acyclic graph (DAG) of dependencies and then creating execution levels where all steps in a level can run in parallel.

Constructors

Constructor

new AxFlowExecutionPlanner(): AxFlowExecutionPlanner;

Returns

AxFlowExecutionPlanner

Methods

addExecutionStep()

addExecutionStep(
   stepFunction: AxFlowStepFunction, 
   nodeName?: string, 
   mapping?: (state: any) => any, 
   stepType?: "map" | "execute" | "merge" | "parallel-map" | "parallel" | "derive", 
   mapTransform?: (state: any) => any, 
   mergeOptions?: object, 
   deriveOptions?: object): void;

Defined in: https://github.com/ax-llm/ax/blob/9a5a7060a48f9eef46efc680b0cdf6b42bff5df2/src/ax/flow/executionPlanner.ts#L53

Adds an execution step to the plan for analysis and optimization.

This method is called for every operation in the flow (execute, map, merge, etc.) and performs dependency analysis to understand what the step needs and produces. This information is crucial for building the parallel execution plan.

The method handles different types of steps:

• Execute steps: LLM node operations that depend on specific state fields
• Map steps: Transformations that modify the state object
• Merge steps: Operations that combine results from branches or parallel operations
• Other steps: Generic operations that don’t fit other categories

Parameters

Returns

void

createOptimizedExecution()

createOptimizedExecution(batchSize?: number): AxFlowStepFunction[];

Defined in: https://github.com/ax-llm/ax/blob/9a5a7060a48f9eef46efc680b0cdf6b42bff5df2/src/ax/flow/executionPlanner.ts#L614

Creates optimized execution functions that implement the parallel execution plan.

This method converts the parallel groups into actual executable functions. It creates a series of steps where:

• Single-step groups execute directly
• Multi-step groups execute in parallel with batch size control
• Results are properly merged to maintain state consistency

The optimized execution can significantly improve performance for flows with independent operations, especially I/O-bound operations like LLM calls.

Performance benefits:

• Reduces total execution time for independent operations
• Maximizes CPU and I/O utilization
• Maintains correctness through dependency management
• Controls resource usage through batch size limiting

Parameters

Returns

AxFlowStepFunction[]

Array of optimized step functions ready for execution

getExecutionPlan()

getExecutionPlan(): object;

Defined in: https://github.com/ax-llm/ax/blob/9a5a7060a48f9eef46efc680b0cdf6b42bff5df2/src/ax/flow/executionPlanner.ts#L704

Gets detailed execution plan information for debugging and analysis.

This method provides comprehensive information about the execution plan, including step counts, parallel grouping details, and the complete dependency structure. It’s particularly useful for:

• Debugging execution flow issues
• Performance analysis and optimization
• Understanding parallelization effectiveness
• Monitoring execution plan complexity

Returns

object

Object containing detailed execution plan metrics and data

getOptimizedExecutionSteps()

getOptimizedExecutionSteps(): AxFlowStepFunction[];

Defined in: https://github.com/ax-llm/ax/blob/9a5a7060a48f9eef46efc680b0cdf6b42bff5df2/src/ax/flow/executionPlanner.ts#L683

Gets optimized execution steps for the flow.

This method provides the optimized execution steps that can be used to execute the flow with maximum parallelism while maintaining dependency order.

Returns

AxFlowStepFunction[]

Array of optimized step functions ready for execution

setInitialFields()

setInitialFields(fields: string[]): void;

Defined in: https://github.com/ax-llm/ax/blob/9a5a7060a48f9eef46efc680b0cdf6b42bff5df2/src/ax/flow/executionPlanner.ts#L475

Sets the initial fields and triggers parallel group rebuilding.

This method is called once the flow knows what input fields are available. It triggers the parallel group analysis which determines the optimal execution strategy for the entire flow.

Parameters

Returns

void