# How to Use AxFlow for Building and Managing AI Workflows in TypeScript

> Learn to build and manage AI workflows in TypeScript with AxFlow. This fluent orchestration engine offers parallel execution, type-safe state, and observability for production-grade AI applications.

- Repository: [Ax/ax](https://github.com/ax-llm/ax)
- Tags: how-to-guide
- Published: 2026-02-25

---

**AxFlow is a fluent, chainable orchestration engine in the Ax framework that lets you declare node graphs with automatic parallel execution, type-safe state management, and production-grade observability.**

AxFlow provides the core workflow infrastructure for the `ax-llm/ax` repository, allowing developers to compose AI programs into sophisticated pipelines. It handles everything from dependency analysis and automatic parallelization to state evolution and error handling, all while maintaining compile-time type safety through TypeScript generics.

## What is AxFlow?

AxFlow serves as the **meta-program execution layer** within the Ax ecosystem. When you invoke `forward()`, the engine performs several critical operations defined in [`src/ax/flow/flow.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/flow/flow.ts):

- **Signature Inference**: Starting with a generic flow created via `flow()`, the engine inspects every node and automatically constructs a precise `AxSignature` that defines the contract of inputs and outputs (see `inferSignatureFromFlow()` at lines 82-106).

- **Dependency Analysis**: The `AxFlowExecutionPlanner` traverses the node graph to determine which steps depend on which fields, grouping independent steps into parallel execution groups (see [`src/ax/flow/executionPlanner.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/flow/executionPlanner.ts) lines 27-73).

- **Automatic Parallelization**: When `autoParallel` is enabled (the default), each parallel group runs concurrently while respecting the configurable `batchSize` (see `forward()` lines 71-80).

- **State Evolution**: A mutable `state` object threads through every step, with results stored under the convention `{nodeName}Result` (see `executeStepsWithLogging()` lines 33-38).

- **Observability**: Optional colorized logs and OpenTelemetry spans emit for each step, providing production-grade tracing (see [`src/ax/flow/logger.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/flow/logger.ts) and `forward()` lines 100-115).

- **Error Handling**: A per-flow abort controller allows `flow.stop()` to cancel in-flight LLM calls, surfacing an `AxAIServiceAbortedError` (see `stop()` lines 46-51).

## Core Concepts and API

Understanding AxFlow's building blocks is essential for constructing efficient workflows.

### Node Registration

Nodes are reusable AI programs defined by signature strings or `AxSignature` objects. Register them using `flow.node(name, signature)` or the compact alias `flow.n()`. According to [`src/ax/flow/flow.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/flow/flow.ts) lines 75-90, each node registration creates an `AxGen` instance that participates in the execution graph.

### Execution

The `flow.execute(name, selector)` method (alias `flow.e()`) runs a specific node, feeding it a subset of the current state. The result automatically becomes available as `{nodeName}Result` in subsequent steps (see [`src/ax/flow/flow.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/flow/flow.ts) lines 124-135).

### State Transformation

Use `flow.map(fn)` (alias `flow.m()`) for synchronous or asynchronous state mutations. This is ideal for preprocessing inputs, post-processing outputs, or adding derived fields without invoking LLM calls (see [`src/ax/flow/flow.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/flow/flow.ts) lines 150-165).

### Conditional Branching

AxFlow supports conditional routing through `branch()`, `when()`, and `merge()`. You define a selector function that determines which branch to follow, execute divergent logic, then join the branches back into a single state using `merge()` (see [`src/ax/flow/flow.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/flow/flow.ts) lines 200-225).

### Parallel Execution

The `parallel()` method accepts an array of sub-flow builders, each receiving its own copy of the state. These execute concurrently, with results combined via a subsequent `merge()` step. This pattern is defined in [`src/ax/flow/flow.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/flow/flow.ts) lines 240-260.

### Looping Constructs

Implement iterative refinement using `while(predicate)` and `endWhile()`. The engine repeats the enclosed steps until the predicate function returns `false`, enabling feedback loops for quality thresholds or convergence criteria (see [`src/ax/flow/flow.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/flow/flow.ts) lines 275-295).

### Factory Methods

Always instantiate flows using the `flow()` factory function or `AxFlow.create()`. These methods inject proper generic types and avoid the deprecated `new AxFlow()` constructor (see [`src/ax/flow/flow.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/flow/flow.ts) lines 55-63).

## Building Your First Workflow

Here is a basic document analysis pipeline that demonstrates node registration, execution, and output extraction:

```typescript
import { ai, flow } from "@ax-llm/ax";

const llm = ai({ name: "openai", apiKey: process.env.OPENAI_APIKEY! });

const docFlow = flow<{ document: string }, { summary: string }>()
  .description("Doc summarizer", "Creates a short summary of a document.")
  .node("summarizer", "document:string -> summary:string")
  .execute("summarizer", (s) => ({ document: s.document }))
  .returns((s) => ({ summary: s.summarizerResult.summary }));

// Run it
const out = await docFlow.forward(llm, { document: "Ax is a TS LLM framework…" });
console.log("Summary:", out.summary);

```

**Key implementation details**:
- `node()` creates an `AxGen` from the signature string
- `execute()` feeds the current state's `document` field to the node
- `returns()` acts as a terminal `map()` that extracts the final output signature

## Advanced Patterns

### Using Aliases for Compact Code

For concise workflow definitions, AxFlow provides single-letter aliases that mirror the full method names:

```typescript
import { flow } from "@ax-llm/ax";

const supportFlow = flow<{ message: string }>()
  .description("Support ticket processor", "Classifies and replies to tickets.")
  .n("classifier", "message:string -> category:string, urgency:string")
  .n("responder", "category:string, urgency:string -> reply:string")
  .e("classifier", (s) => ({ message: s.message }))
  .e("responder", (s) => ({
    category: s.classifierResult.category,
    urgency: s.classifierResult.urgency,
  }))
  .m((s) => ({ reply: s.responderResult.reply }));

const reply = await supportFlow.forward(aiInstance, { message: "My order is broken" });
console.log(reply.reply);

```

The **alias methods** (`n`, `e`, `m`) reduce boilerplate while maintaining identical functionality to their full-name counterparts.

### Parallel Processing and Iteration

This research paper scorer demonstrates parallel evaluation with iterative refinement:

```typescript
import { flow } from "@ax-llm/ax";

const scorer = flow<{ abstract: string }>()
  .description("Research paper scorer", "Scores novelty & clarity in parallel, iterates until stable.")
  .node("novelty", "abstract:string -> novelty:number")
  .node("clarity", "abstract:string -> clarity:number")
  .parallel([
    (sub) => sub.execute("novelty", (s) => ({ abstract: s.abstract })),
    (sub) => sub.execute("clarity", (s) => ({ abstract: s.abstract })),
  ])
  .merge("combinedScore", (noveltyRes, clarityRes) => {
    const n = (noveltyRes as any).noveltyResult.novelty;
    const c = (clarityRes as any).clarityResult.clarity;
    return (n + c) / 2;
  })
  .while((s) => s.combinedScore < 0.8)
  .map((s) => ({
    abstract: s.abstract + " (refined)",
    combinedScore: s.combinedScore + 0.1,
  }))
  .endWhile()
  .returns((s) => ({ finalScore: s.combinedScore }));

const result = await scorer.forward(aiInstance, { abstract: "..." });
console.log("Final score:", result.finalScore);

```

The `parallel()` array executes both scoring nodes concurrently. The `while()` loop continues refinement until the `combinedScore` meets the quality threshold.

## Execution Engine and Performance

AxFlow optimizes workflow execution through sophisticated planning. The `AxFlowExecutionPlanner` in [`src/ax/flow/executionPlanner.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/flow/executionPlanner.ts) constructs a dependency graph that identifies which nodes can run simultaneously based on their input requirements.

When `autoParallel` is enabled (the default setting), the engine automatically batches independent operations. You can control concurrency through the `batchSize` parameter, limiting how many nodes execute simultaneously to prevent rate limiting or resource exhaustion.

The engine maintains type safety throughout execution by updating the generic `TState` type with each chained method call. This provides IntelliSense for state fields like `summarizerResult` or `classifierResult` immediately after defining the corresponding nodes.

## Observability and Error Handling

AxFlow integrates with the broader Ax ecosystem for monitoring and control. It reuses the global `axGlobals.cachingFunction` and OpenTelemetry tracer defined in `forward()` lines 111-120, providing end-to-end visibility across workflow execution.

For long-running workflows, the abort controller pattern allows graceful cancellation. Calling `flow.stop()` triggers the abort signal, causing pending LLM calls to surface `AxAIServiceAbortedError` rather than hanging indefinitely (see `stop()` lines 46-51).

## Summary

- **AxFlow** is the orchestration engine in `ax-llm/ax` that manages AI program execution through a fluent, chainable API.
- Use the `flow()` factory to create type-safe workflows with methods like `node()`, `execute()`, `map()`, `branch()`, `parallel()`, and `while()`.
- The **execution planner** automatically parallelizes independent steps when `autoParallel` is enabled, respecting your configured `batchSize`.
- State follows the `{nodeName}Result` naming convention, with full TypeScript inference tracking available fields throughout the chain.
- Production features include **OpenTelemetry tracing**, **colorized logging**, and **abort controller** support for cancellation.

## Frequently Asked Questions

### What is the difference between AxFlow and AxProgram?

**AxProgram** is the low-level execution primitive that individual nodes wrap, while **AxFlow** builds a meta-program that orchestrates multiple nodes. According to the source code in [`src/ax/flow/flow.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/flow/flow.ts) lines 73-79, AxFlow registers every node with a shared `AxProgram` instance via `ensureProgram()`, allowing the framework to manage dependencies and execution order across the entire workflow graph.

### How does AxFlow handle parallel execution?

AxFlow analyzes dependencies through the `AxFlowExecutionPlanner` to determine which nodes can run simultaneously. When `autoParallel` is enabled (default), independent nodes execute concurrently in grouped batches. You can manually define parallel blocks using `parallel()` followed by `merge()` to combine results, as implemented in [`src/ax/flow/flow.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/flow/flow.ts) lines 240-260.

### Can I stop a running workflow in AxFlow?

Yes. Each flow instance maintains an abort controller that you can trigger via `flow.stop()`. This cancels any in-flight LLM calls and surfaces an `AxAIServiceAbortedError` to the caller. This pattern is essential for production applications where user cancellation or timeout handling is required (see [`src/ax/flow/flow.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/flow/flow.ts) lines 46-51).

### Is AxFlow type-safe?

AxFlow provides full TypeScript type safety through generic state tracking. Each call to `node()`, `execute()`, or `map()` updates the `TState` generic type, ensuring compile-time IntelliSense for state fields like `classifierResult` or `combinedScore`. The factory method `flow()` properly initializes these generics, avoiding the type inference issues present in the deprecated `new AxFlow()` constructor.