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

> Implement a custom LLM provider for Ax by extending AxAIService and wrapping it with AxAIProvider. Integrate any language model seamlessly with the Ax SDK for full compatibility. Get the complete guide.

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

---

**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`](https://github.com/ax-llm/ax/blob/main/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`](https://github.com/ax-llm/ax/blob/main/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`](https://github.com/ax-llm/ax/blob/main/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`](https://github.com/ax-llm/ax/blob/main/src/ax/ai/mock/api.ts) and the patterns shown in [`src/docs/src/content/docs/quickstart.md`](https://github.com/ax-llm/ax/blob/main/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:

```typescript
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:

```typescript
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`](https://github.com/ax-llm/ax/blob/main/src/docs/src/content/docs/quickstart.md) (lines 319–338):

```typescript
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`](https://github.com/ax-llm/ax/blob/main/src/ax/ai/types.ts)** – Defines the `AxAIService` contract that every provider must implement (lines 889–916)
- **[`src/aisdk/provider.ts`](https://github.com/ax-llm/ax/blob/main/src/aisdk/provider.ts)** – Contains `AxAIProvider` which bridges `AxAIService` to the Vercel AI SDK (lines 129–160)
- **[`src/ax/ai/mock/api.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/ai/mock/api.ts)** – Reference implementation `AxMockAIService` serving as a minimal template (lines 51–66)
- **[`src/docs/src/content/docs/quickstart.md`](https://github.com/ax-llm/ax/blob/main/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`](https://github.com/ax-llm/ax/blob/main/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`](https://github.com/ax-llm/ax/blob/main/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`](https://github.com/ax-llm/ax/blob/main/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.