DOCSEARCHbyAlgolia
Submit Docs

Ax Documentation

repository·main·Indexed 25 days ago

https://github.com/ax-llm/ax

Ax is a DSPy-inspired multi-language programming model for building reliable AI applications. It provides a unified semantic core for defining typed signatures, agents, and flows across TypeScript, Python, Java, C++, Go, and Rust. Key features include AxAI for provider normalization (OpenAI, Gemini, Anthropic, etc.), AxGen for structured generation, AxFlow for program graph composition, and AxAgent for RLM actor-code execution with optional runtime profiles like QuickJS and Pyodide.

Tokens
143.4K
Snippets
398
Records
767
Agent score
83%

What's inside Ax

  1. Overview of Ax for Python capabilities

    main

    Ax for Python allows you to build reliable AI applications using a shared compiler contract with other languages. Key features include:

    • Signatures and schemas: Define inputs/outputs once for reuse in prompts, tools, and validation.
    • AxGen: Structured generation with retries, tool calls, field processors, assertions, and provider-backed parsing.
    • AxAI: Unified provider boundary for OpenAI-compatible APIs, Gemini, Anthropic, Azure OpenAI, DeepSeek, Mistral, Reka, Cohere, and Grok.
    • Audio and realtime: Support for .chat() with input_audio, transcribe()/speak() for batch processing, and realtime_chat() for streaming audio over WebSockets.
    • AxAgent and RLM: Agentic planning and execution using an RLM executor loop that manages state, logs, and traces.
    • AxFlow: Composition of AxGen, AxAgent, and nested flows into portable program graphs.
    • Optimizers: Lifecycle management (save, load, apply, evaluate) for optimizer artifacts like the GEPA engine.

  2. Overview of Ax Runtime Surfaces

    main

    Ax provides seven main runtime surfaces for building AI applications:

    1. AxAI: Handles provider clients, model metadata, chat, streaming, embeddings, audio (transcribe/speak), and routing.
    2. AxGen: Provides signature-driven structured generation, including prompts, tools, retries, schema validation, assertions, bestOfN(...), and refine(...).
    3. AxAgent: A staged agent pipeline supporting actor runtime sessions, delegation, context budgets, and state export/restore.
    4. AxFlow: An Ax program graph supporting dependency planning, auto-parallel grouping, control flow (branch/while/feedback), and caching.
    5. AxOptimize: Manages optimizable components, evaluator rollouts, and optimizer engines like GEPA.
    6. AxEventRuntime: A protocol-neutral layer for observing, routing, and resuming Ax programs via a durable-inbox.
    7. AxIR generated libraries: Language-agnostic libraries (Python, Java, C++, Go, Rust) emitted from shared portable semantics.

  3. Overview of Ax capabilities in Java

    main

    Ax provides a native Java API for building reliable AI applications with the following features:

    • Signatures and Schemas: Define input/output shapes once and reuse them for validation and typed results.
    • AxGen: Perform structured generation with retries, tool calls, field processors, assertions, and provider-backed parsing.
    • AxAI: Unified provider boundary for OpenAI-compatible APIs, Gemini, Anthropic, Azure OpenAI, DeepSeek, Mistral, Reka, Cohere, and Grok.
    • Audio and Realtime: Support for .chat() with input_audio, batch speech-to-text/text-to-speech via .transcribe()/.speak(), and WebSocket-based realtime audio.
    • AxAgent and RLM: Agentic planning and execution using an RLM executor loop that manages state, logs, traces, and context.
    • AxFlow: Compose AxGen, AxAgent, and nested flows into portable program graphs.
    • Optimizers: Save, load, and apply optimizer artifacts (e.g., GEPA engine).

  4. Understand the AxJSRuntime threat model

    main

    When using AxJSRuntime for executing LLM-authored code, treat it as a defense-in-depth mechanism rather than a secure container or VM boundary.

    Security Best Practices:

    • The host callbacks and granted runtime permissions are the actual authority boundaries.
    • Never store durable secrets or privileged effects directly within the LLM-authored code or the runtime environment.
    • Always keep secrets and sensitive operations in host-side functions that are explicitly granted to the runtime.

  5. Understand the AxIR Compiler and Portability Layer

    main
    AxIR is the compiler-owned portability layer for Ax. It translates the shared Ax runtime contract into native libraries for Python, Java, C++, Go, and Rust. While TypeScript serves as the behavioral reference implementation, AxIR ensures that core Ax semantics (like signature parsing, tool-call normalization, and event routing) remain deterministic and language-agnostic across all supported languages.

  6. Understand the AxIR Lowering Pipeline

    main

    The Ax compiler uses a staged lowering process to transform .axir modules into executable code. The pipeline follows these stages:

    1. parse: Reads .axir modules.
    2. resolve: Loads imports and collects symbols.
    3. check: Validates dialect names, references, public symbols, and Core body invariants.
    4. canonicalize: Produces stable formatter output and operation ordering.
    5. lower --to core: Converts Ax dialect operations with core_kind into Core operations.
    6. core-body: Validates and normalizes typed Core regions for executable semantic bodies.
    7. runtime-model: Extracts the Ax runtime package model and target idiom contract from lowered Core.
    8. compile: Emits source from the runtime model.

    This process ensures that portable executable semantics (Core bodies) are separated from host-specific implementations (Target boundaries).

  7. Understand the Ax Program Dialect

    main
    The ax.program dialect defines the shared contract for all executable Ax programs. It provides a consistent interface for cross-program semantics, including descriptors, trace events, component metadata, and rollout hooks. This thin dialect ensures that different program types can interact and be traced using a unified shape.

  8. AxGen Java Package Facts

    main

    The ax-java-gen skill provides the following capabilities for Java development:

    • Package Name: dev.axllm:ax
    • Language: Java
    • Network Support: Real network support is available.
    • Transport: Supports scripted no-key transport.
    • Runtime Profiles: javascript-quickjs, python-pyodide.
    • Documentation Sources: API.md, axir-api.json, and axir-capabilities.json.

  9. Understand AxIR Provider Parity and Boundaries

    main

    AxIR (Ax Intermediate Representation) ensures semantic portability across different AI providers. When using Ax, it is important to distinguish between what is handled by the Ax core (semantics, mapping, and conformance) and what is handled by the host/transport layer.

    Core AxIR Responsibilities (Handled by Ax)

    • Request/Response Mapping: Normalizing inputs and outputs across providers.
    • Stream Folding: Consolidating stream events.
    • Usage Normalization: Standardizing token/cost usage data.
    • Descriptor Identity: Managing provider aliases and model catalogs.
    • Generated Conformance: Ensuring providers like Anthropic, Gemini, and OpenAI-compatible services follow AxIR semantics.
    • Routing & Balancing: Logic for AxProviderRouter, AxMultiServiceRouter, and AxBalancer (including scoring, selection, and failover).

    Host/Transport Boundaries (NOT handled by Ax)

    The following tasks are the responsibility of the host implementation or the specific provider's live transport layer:

    • Live Network Execution: HTTP, SSE, WebSocket, and multipart transports.
    • Authentication: Loading and managing auth credentials.
    • Retries: Timer-backed live backoff and product-level retry policies.
    • Media Handling: Binary uploads and media-device handling.
    • Realtime Sessions: Provider-side realtime session management.

  10. Ax Go Audio Package Reference

    main

    The ax-go-audio skill provides capabilities for audio input/output, OpenAI Responses audio mapping, and realtime event folding using the Go package github.com/ax-llm/ax/packages/go.

    Key Capabilities

    • Mapping speech, transcription, or realtime events through the generated provider surface.
    • Event folding and provider request mapping using no-key examples.
    • Real network support and scripted no-key transport support.
    • Runtime profiles: javascript-goja.

    Core API Surface

    • Clients: NewAI, OpenAICompatibleClient, OpenAIResponsesClient, GoogleGeminiClient, AnthropicClient.
    • Usage & Observability: AxUsageContext, AxUsageEvent, AxUsageObserver, SetUsageObserver.
    • Routing & Balancing: AxBalancer, AxBalancerAdaptiveStrategy, AxBalancerStatsStore, AxInMemoryBalancerStatsStore, CreateBalancerRouteStats, UpdateBalancerRouteStats, SampleBalancerRouteHealth, MultiServiceRouter, ProviderRouter.

  11. Core Ax Concepts

    main

    Ax provides a unified programming model for LLM applications through several core primitives:

    • Signatures: Define typed structured generation using a string DSL, a fluent f() builder, or Standard Schema v1 validators (Zod, Valibot, ArkType).
    • Provider Abstraction: A single interface for OpenAI-compatible endpoints, Anthropic, Gemini, Grok/xAI, Mistral, Cohere, Reka, DeepSeek, Azure OpenAI, and more.
    • Agents: Runtime execution units with context budgets, checkpoints, action-log replay, memory, and delegation.
    • Flows: Typed program graphs supporting branches, loops, feedback, caching, and parallel execution.
    • Optimizers: Tools for optimization including GEPA and few-shot bootstrapping.

  12. Understand the AxFlow Dialect

    main

    AxFlow models an Ax program as a directed graph. Unlike generic workflow engines, AxFlow is specifically designed for Ax programs where every executable node represents either an Ax program boundary or a deterministic host callback.

    Key features of the AxFlow alpha contract include:

    • Graph construction and duplicate step checks.
    • Planning metadata and cache-key generation derived from actual input objects.
    • Support for child-program forward calls.
    • Trace and chat-log aggregation.
    • .returns() projection for extracting specific outputs from the flow.