How to Add a Custom LLM Provider to Ax: A Complete Implementation Guide

Ax abstracts LLM access through the AxAIService interface, allowing you to integrate any language model by implementing this contract and wrapping it with AxAIProvider for full SDK compatibility.

Adding a custom LLM provider to the Ax framework (ax-llm/ax) requires implementing the AxAIService interface defined in the source code. This pattern enables you to connect proprietary APIs, experimental models, or local inference engines to Ax's ecosystem. Once implemented, the AxAIProvider wrapper bridges your service to the Vercel AI SDK, unlocking flows, DSP orchestration, caching, and telemetry.

Understanding the AxAIService Interface

The foundation of any custom provider starts with the contract defined in src/ax/ai/types.ts. The AxAIService interface specifies the methods Ax calls for chat completion, embeddings, model metadata, metrics, and option handling.

Core Contract Requirements

The interface requires implementation of several key methods that define your LLM's identity and capabilities:

• getId() and getName() – Return unique identifiers for the service
• getFeatures() – Declare supported capabilities (streaming, function calling, media, etc.)
• chat() – Handle text generation requests and return responses or streams
• embed() – Provide vector embeddings (optional but recommended)
• getModelList() – Return available models (can return undefined)
• getLastUsedModelConfig() – Return the configuration of the last used model
• getMetrics() and getLogger() – Support observability and telemetry
• setOptions() and getOptions() – Manage runtime configuration

As implemented in src/ax/ai/types.ts (lines 889–916), this contract ensures Ax can uniformly interact with diverse LLM backends.

Bridging to the Vercel AI SDK with AxAIProvider

Located in src/aisdk/provider.ts, the AxAIProvider class acts as a translation layer between your AxAIService implementation and the Vercel AI SDK's LanguageModelV2 interface. According to the source code (lines 129–160), this wrapper translates Ax-style calls into the SDK's doGenerate and doStream methods.

This bridge allows your custom provider to work anywhere the Vercel AI SDK is expected while maintaining access to Ax's higher-level features like flow orchestration and DSP (Digital Signal Processing) patterns.

Step-by-Step Custom Provider Implementation

Follow this three-step process to integrate your LLM, based on the reference implementation in src/ax/ai/mock/api.ts and the patterns shown in src/docs/src/content/docs/quickstart.md.

Step 1: Implement the AxAIService Interface

Create a class that implements the required methods. This example shows a minimal implementation that echoes requests:

import type {
  AxAIService,
  AxChatRequest,
  AxChatResponse,
  AxEmbedRequest,
  AxEmbedResponse,
  AxAIServiceMetrics,
  AxModelConfig,
} from '@ax-llm/ax';

class MyCustomLLM implements AxAIService {
  // Basic identity
  getId() { return 'my-llm-id'; }
  getName() { return 'my-llm'; }

  // Feature flagging (adjust to your LLM's capabilities)
  getFeatures() {
    return {
      functions: true,
      streaming: true,
      structuredOutputs: false,
      media: { 
        images: { supported: false, formats: [] }, 
        audio: { supported: false, formats: [] }, 
        files: { supported: false, formats: [], uploadMethod: 'none' }, 
        urls: { supported: false, webSearch: false, contextFetching: false } 
      },
      caching: { supported: false, types: [] },
      thinking: false,
      multiTurn: true,
    };
  }

  // Optional model list / config
  getModelList() { return undefined; }
  
  getLastUsedModelConfig(): AxModelConfig | undefined {
    return { maxTokens: 1024, temperature: 0.7, stream: true };
  }

  // Metrics (simple stub)
  private metrics: AxAIServiceMetrics = { 
    latency: { 
      chat: { mean: 0, p95: 0, p99: 0, samples: [] }, 
      embed: { mean: 0, p95: 0, p99: 0, samples: [] } 
    }, 
    errors: { 
      chat: { count: 0, rate: 0, total: 0 }, 
      embed: { count: 0, rate: 0, total: 0 } 
    } 
  };
  
  getMetrics() { return this.metrics; }

  // Logger placeholder
  getLogger() { return () => {}; }

  // Required chat implementation
  async chat(
    req: Readonly<AxChatRequest<unknown>>
  ): Promise<AxChatResponse | ReadableStream<AxChatResponse>> {
    // Translate Ax request → your LLM API request, then map response back
    const response: AxChatResponse = {
      results: [{ 
        index: 0, 
        content: `Echo: ${JSON.stringify(req)}`, 
        finishReason: 'stop' 
      }],
      modelUsage: { 
        ai: this.getName(), 
        model: 'custom-model', 
        tokens: { 
          promptTokens: 5, 
          completionTokens: 5, 
          totalTokens: 10 
        } 
      },
    };
    return response;
  }

  // Optional embed implementation
  async embed(req: Readonly<AxEmbedRequest>): Promise<AxEmbedResponse> {
    return { 
      vectors: [], 
      modelUsage: { 
        ai: this.getName(), 
        model: 'embed-model', 
        tokens: { 
          promptTokens: 0, 
          completionTokens: 0, 
          totalTokens: 0 
        } 
      } 
    };
  }

  // Option handling (store for later calls)
  private opts: Readonly<any> = {};
  
  setOptions(options: any) { this.opts = options; }
  getOptions() { return this.opts; }
}

Step 2: Wrap with AxAIProvider

Import AxAIProvider from the AI SDK package and instantiate it with your custom service:

import { AxAIProvider } from '@ax-llm/ax-ai-sdk-provider';

const myProvider = new AxAIProvider(new MyCustomLLM());

This wrapper implements the LanguageModelV2 interface, making your provider compatible with Vercel AI SDK conventions.

Step 3: Integrate with Ax Workflows

Pass the wrapped provider directly to Ax flows, agents, or DSP functions. As shown in src/docs/src/content/docs/quickstart.md (lines 319–338):

import { ax } from '@ax-llm/ax';

const generator = ax(`
  prompt:string "Your prompt"
  -> result:string "LLM output"
`);

await generator.forward(myProvider, { prompt: 'Hello world' });

You can also use the provider with Ax's ai() factory or custom agents by passing the instance where an AI service is expected.

Reference Implementation and Key Files

The repository provides several critical files for implementing custom providers:

• src/ax/ai/types.ts – Defines the AxAIService contract that every provider must implement (lines 889–916)
• src/aisdk/provider.ts – Contains AxAIProvider which bridges AxAIService to the Vercel AI SDK (lines 129–160)
• src/ax/ai/mock/api.ts – Reference implementation AxMockAIService serving as a minimal template (lines 51–66)
• src/docs/src/content/docs/quickstart.md – Documentation showing instantiation and integration patterns (lines 319–338)

Use AxMockAIService as a starting point when building production providers, as it demonstrates the minimal required surface area of the interface.

Summary

• Implement AxAIService from src/ax/ai/types.ts to define your LLM's identity, capabilities, and generation logic
• Wrap your implementation with AxAIProvider to achieve Vercel AI SDK compatibility and access standardized methods
• Configure feature flags in getFeatures() to accurately reflect whether your model supports streaming, function calling, structured outputs, or media
• Reference src/ax/ai/mock/api.ts as a minimal working template for new providers
• Pass the wrapped provider directly to Ax flows, agents, and DSP generators using the .forward() method or factory functions

Frequently Asked Questions

What methods are required to implement AxAIService?

You must implement getId(), getName(), getFeatures(), chat(), getMetrics(), getLogger(), setOptions(), and getOptions(). The embed() method is optional if your provider doesn't support embeddings, and getModelList() can return undefined if your LLM doesn't expose a model catalog. Each method serves a specific purpose in the Ax ecosystem, from capability declaration to actual text generation and observability.

Do I need to implement the embed method for my custom LLM provider?

No, the embed() method is optional. If your LLM doesn't support text embeddings, you can either omit the method or provide a stub that throws an unsupported operation error. However, if you plan to use Ax's vector search, RAG features, or any embedding-dependent DSP patterns, you'll need a functional implementation that returns AxEmbedResponse with embedding vectors.

How does AxAIProvider bridge my custom service to the Vercel AI SDK?

According to src/aisdk/provider.ts, AxAIProvider receives your AxAIService instance and implements the LanguageModelV2 interface expected by the Vercel AI SDK. It translates Ax's chat() calls into the SDK's doGenerate() and doStream() methods, allowing your custom LLM to work with any Vercel AI SDK compatible tooling while maintaining access to Ax's higher-level features like flow orchestration and caching.

Can I use feature flags to control my custom provider's capabilities?

Yes, the getFeatures() method returns a configuration object that tells Ax what your LLM supports. Set boolean values for functions, streaming, structuredOutputs, thinking, and multiTurn, and configure media handling for images, audio, and files. Accurate feature declaration ensures Ax only sends compatible requests—such as disabling function calls for models that don't support tool use or enabling streaming only when your backend supports server-sent events.

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