How Streaming Output Works in Ax: Implementation and Handling Guide

Ax implements streaming output as an asynchronous generator that yields incremental deltas while processing a ReadableStream, allowing real-time consumption of partial LLM responses through the streamingForward() method.

Streaming output in Ax enables real-time processing of large language model responses before they complete. The ax-llm/ax repository provides a robust streaming architecture through its DSP (Declarative Signature Programming) module, handling everything from chunk parsing to validation. This guide examines the internal mechanisms and practical patterns for consuming streaming output in production applications.

The Streaming Architecture

Ax breaks streaming into three distinct stages: request creation, chunk processing, and finalization. Each stage operates asynchronously to transform raw byte streams into structured, incremental outputs.

Request Initialization with streamingForward

The entry point for streaming lives in src/ax/dsp/generate.ts. The AxGen.streamingForward() method prepares the generation context and initiates the network request:

public async *streamingForward<T extends Readonly<AxAIService>>(
  ai: T,
  values: IN | AxMessage<IN>[],
  options?: Readonly<AxProgramForwardOptions<any>>
): AxGenStreamingOut<OUT> { … }

This method resets the prompt-template manager, resolves stop-function names, and prepares mutable options including functions and self-tuning parameters. It delegates to forwardCore, which eventually reaches forwardSendRequest. When the AI service returns a ReadableStream, control passes to processStreamingResponse.

Chunk Processing in processStreamingResponse

The core streaming logic resides in src/ax/dsp/processResponse.ts within the processStreamingResponse function:

export async function* processStreamingResponse<OUT extends AxGenOut>({
  res,
  usage,
  states,
  debug,
  stepContext,
  …args
}: ProcessStreamingResponseArgs): AsyncGenDeltaOut<OUT> { … }

The function implements a read loop that consumes res.getReader().read() until the stream signals completion. For each chunk, it performs several operations:

• Usage and citation extraction – Accumulates modelUsage metadata and citation data across chunks
• Empty-chunk filtering – Skips chunks lacking content, thoughts, or function calls
• State matching – Locates the InternalAxGenState corresponding to result.index
• Delta computation – Transforms raw chunks into effective deltas containing only newly-available data

The processor handles incomplete JSON structures through parsePartialJson (located in src/ax/util/partialJson.ts), enabling robust parsing of partial schema outputs. It then executes assertions via assertAssertions, streaming assertions via assertStreamingAssertions, and applies field processors through processFieldProcessors and processStreamingFieldProcessors before yielding the delta.

Stream Finalization

After the read loop terminates, finalizeStreamingResponse flushes any remaining buffered data for each state. This ensures that JSON objects split across multiple chunks or trailing content gets properly emitted before the generator completes.

Consuming Streaming Output

AxGen.streamingForward returns an async iterator that yields structured deltas progressively. This pattern enables UI updates, logging, and pipeline processing without waiting for full response completion.

Basic Async Iterator Pattern

Consume the stream using a standard for await...of loop:

import { AxAI, AxAIGoogleGeminiModel, ax } from '@ax-llm/ax';

const gen = ax(
  'movieTitle:string -> rating:number, genres:string[], verdict:string'
);

const ai = new AxAI({
  name: 'google-gemini',
  apiKey: process.env.GOOGLE_APIKEY!,
  config: { model: AxAIGoogleGeminiModel.Gemini20FlashLite },
});

const generator = await gen.streamingForward(ai, { movieTitle: 'The Grand Budapest Hotel' });

for await (const delta of generator) {
  // Each delta contains only new fields or updated values
  console.log('Partial output:', delta);
}

This example from src/examples/streaming.ts demonstrates the standard consumption pattern where each iteration receives an object representing the incremental progress toward the final structured output.

Handling Deltas and Partial Updates

Each yielded value represents an effective delta—only the fields that changed or newly appeared in that specific chunk. When building user interfaces, merge these deltas into existing state rather than replacing it:

let accumulatedOutput = {};

for await (const delta of generator) {
  accumulatedOutput = { ...accumulatedOutput, ...delta };
  renderUpdate(accumulatedOutput);
}

This approach preserves partial data across iterations, crucial for handling fields that arrive in arbitrary order or across multiple chunks.

Advanced Streaming Features

Ax provides mechanisms for real-time validation and lifecycle management during streaming operations.

Real-Time Validation with Streaming Assertions

Validate partial output as it arrives using addStreamingAssert. These assertions evaluate on every chunk that updates the target field:

gen.addStreamingAssert('verdict', ({ verdict }) => {
  if (verdict && verdict.length < 50) {
    return `Verdict too short (${verdict.length}); need ≥ 50 characters.`;
  }
  return true;
});

If the assertion fails, Ax immediately aborts generation and surfaces the error. This capability, implemented in src/ax/dsp/asserts.ts, prevents downstream processing of invalid partial outputs.

Cancellation and Error Handling

Supply an AbortSignal through the options parameter to enable cancellation:

const controller = new AbortController();
setTimeout(() => controller.abort('User cancelled'), 5000);

const generator = await gen.streamingForward(ai, { movieTitle: 'Inception' }, {
  abortSignal: controller.signal,
});

try {
  for await (const delta of generator) {
    console.log(delta);
  }
} catch (e) {
  if (e instanceof AxAIServiceAbortedError) {
    console.warn('Streaming aborted:', e.message);
  } else {
    throw e;
  }
}

The read loop checks abortSignal?.aborted after each chunk iteration, throwing AxAIServiceAbortedError when cancelled. Additionally, AxStopFunctionCallException signals graceful termination when stop-functions trigger completion.

Summary

• Ax streams via async generators yielding incremental deltas from streamingForward() in src/ax/dsp/generate.ts.
• Chunk processing occurs in processStreamingResponse (src/ax/dsp/processResponse.ts), handling ReadableStream parsing, partial JSON extraction, and field processing.
• Effective deltas contain only newly-available data, enabling efficient UI updates and pipeline processing.
• Streaming assertions allow real-time validation of partial outputs, aborting immediately on validation failures.
• Cancellation support uses standard AbortController patterns, checked after each chunk to ensure prompt resource cleanup.

Frequently Asked Questions

How do I start a streaming request in Ax?

Call streamingForward() on an AxGen instance, passing an AI service and input values. The method returns an async generator that you consume with for await...of. Unlike the standard forward() method, streamingForward() immediately returns control and begins yielding partial results as they arrive from the LLM provider.

What format does the streaming output use?

Each yield returns an object representing the effective delta of structured output. Rather than streaming raw text tokens, Ax parses chunks into partial JSON objects that conform to your signature's output schema. The parsePartialJson utility in src/ax/util/partialJson.ts handles incomplete structures, ensuring valid partial objects even when JSON spans multiple network chunks.

Can I validate content while it streams?

Yes. Use gen.addStreamingAssert(fieldName, validator) to register functions that evaluate partial field values on every chunk. These streaming assertions run in src/ax/dsp/asserts.ts and can abort generation immediately when validation fails, preventing wasted tokens on outputs that violate your constraints.

How do I cancel an ongoing stream?

Pass an AbortSignal through the options.abortSignal parameter when calling streamingForward(). Ax checks this signal after processing each chunk in the read loop. When aborted, the generator throws AxAIServiceAbortedError, allowing you to distinguish cancellation from other errors and clean up resources appropriately.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}