Ax AI

AI Provider Codegen Rules (@ax-llm/ax)

Use this skill to generate AI provider setup, configuration, and chat code. Prefer short, modern, copyable patterns. Do not write tutorial prose unless the user explicitly asks for explanation.

Quick Setup

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

const openai = ai({ name: 'openai', apiKey: 'sk-...' });
const claude = ai({ name: 'anthropic', apiKey: 'sk-ant-...' });
const gemini = ai({ name: 'google-gemini', apiKey: 'AIza...' });
const azure = ai({ name: 'azure-openai', apiKey: 'your-key', resourceName: 'your-resource', deploymentName: 'gpt-5-4-mini' });
const deepseek = ai({ name: 'deepseek', apiKey: 'sk-...' });
const mistral = ai({ name: 'mistral', apiKey: 'your-key' });
const cohere = ai({ name: 'cohere', apiKey: 'your-key' });
const custom = ai({
  name: 'openai-compatible',
  apiKey: process.env.PROVIDER_API_KEY,
  apiURL: 'https://example.com/v1',
  config: { model: 'provider/model-name' },
});
const reka = ai({ name: 'reka', apiKey: 'your-key' });
const grok = ai({ name: 'grok', apiKey: 'your-key' });

name selects a deployment profile; config.model selects a model inside that deployment. Never infer provider behavior from a model ID. For example, name: 'together' with a deepseek-ai/... model uses Together's profile rules, not DeepSeek's native request shape.

Use axAIProfiles() to discover the complete named catalog and axGetAIProfile(name) to inspect endpoint requirements, authentication, operations, capabilities, model rules, sources, and review dates. Use name: 'openai-compatible' plus apiURL for an unlisted custom endpoint; unknown names are errors.

Profile-only branded classes were removed in the major-version migration. Use ai({ name: ... }) for Azure OpenAI, Cohere, DeepSeek, DeepSeek Responses, Mistral, Reka, and Grok. Retained low-level classes represent genuine transports/runtimes only; legacy model enum and catalog exports remain usable.

Renewable Credentials

Use credentialProvider for expiring bearer tokens. The callback runs for each request attempt and receives { profile, operation, method, url }. Fresh headers override static authentication headers.

const vertex = ai({
  name: 'vertex-ai',
  apiURL: process.env.VERTEX_AI_API_URL!,
  config: { model: 'google/gemma-4-26b-a4b-it-maas' },
  credentialProvider: async ({ operation, url }) => ({
    Authorization: `Bearer ${await tokenSource.fresh({ operation, url })}`,
  }),
});

• Authentication-required profiles accept apiKey or credentialProvider.
• The hook covers chat, stream, embeddings, Responses, transcription, speech, and retry attempts. A hook error stops before transport.
• Do not expect an automatic replay after a completed 401 or 403; generation requests are not safely idempotent.
• Keep ADC or cloud-SDK dependencies in the host. Ax core only owns the neutral callback contract.

Vertex capability rules are model-aware. Documented Gemini MaaS IDs prefer native schema. The exact google/gemma-4-26b-a4b-it-maas rule prefers json_object, excludes native schema, defaults thinking to max, writes chat_template_kwargs.enable_thinking, and extracts/replays reasoning_content. Unknown Vertex models remain conservative.

<!-- axir-nonportable:start webllm -->

WebLLM is browser-only and requires a host-created WebLLM engine. The host loads or reloads models with WebLLM APIs such as CreateMLCEngine(...); Ax only forwards chat requests to that loaded engine. Do not present WebLLM as a portable AxIR provider or a server-side default.

import { ai, AxAIWebLLMModel } from '@ax-llm/ax';

const engine = await CreateMLCEngine(AxAIWebLLMModel.Llama32_3B_Instruct);
const llm = ai({
  name: 'webllm',
  engine,
  config: {
    model: AxAIWebLLMModel.Llama32_3B_Instruct,
    stream: false,
    supportsFunctions: false,
  },
});

<!-- axir-nonportable:end webllm -->

Model Presets

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

const gemini = ai({
  name: 'google-gemini',
  apiKey: process.env.GOOGLE_APIKEY!,
  config: { model: 'simple' },
  models: [
    { key: 'tiny', model: AxAIGoogleGeminiModel.Gemini35FlashLite, description: 'Fast + cheap', config: { maxTokens: 1024 } },
    { key: 'simple', model: AxAIGoogleGeminiModel.Gemini37Flash, description: 'Balanced' },
  ],
});

await gemini.chat({ model: 'tiny', chatPrompt: [{ role: 'user', content: 'Hi' }] });

Model Catalog

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

const providers = axGetSupportedAIModels();
const openai = providers.find((provider) => provider.name === 'openai');
console.log(openai?.models[0]?.promptTokenCostPer1M);
console.log(openai?.capabilities.serviceTiers);

const reasoningModel = openai?.models.find(
  (model) => model.capabilities.thinkingLevels.includes('high')
);
console.log(reasoningModel?.capabilities.thinkingLevels);

const textProviders = axGetSupportedAIModels({ type: 'text' });
const embeddingProviders = axGetSupportedAIModels({ type: 'embeddings' });

Use axGetSupportedAIModels() to build provider/model selectors before creating an ai(...) instance. It returns bundled static metadata: provider names, display names, default models, raw AxModelInfo pricing/details, model type ('text', 'embeddings', 'code', or 'audio'), and normalized capabilities for thinking, thoughts, structured outputs, audio, temperature, top-p, portable thinking levels, and verified service tiers. Provider capabilities describe the default deployment profile; each static model also carries its resolved capabilities.thinkingLevels and capabilities.serviceTiers. Portable thinking levels can collapse onto the same provider-native value. serviceTiers lists verified explicit tiers; serviceTier: 'auto' remains available as the provider-delegated policy.

Provider groups and models are sorted cheapest to most expensive based on bundled input + output token pricing; unpriced models sort last. Dynamic profiles remain useful even when models is empty because their provider-level capability metadata is still returned.

Filter with { type: 'all' | 'text' | 'embeddings' | 'code' | 'audio' } or an array of those values. The 'text' filter includes code-capable models; use 'code' to show only code-first models.

Dynamic providers such as Azure OpenAI deployments are marked with isDynamic: true and may have an empty or static-limited model list.

Portable Inference Service Tiers

Use the shared serviceTier option with auto, standard, flex, or priority. Set it per call, as an instance option, or on a model-key preset; the narrower setting wins. auto delegates to the provider and is omitted when that provider has no explicit auto value. Ax rejects unsupported explicit tiers before transport.

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

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

const response = await gemini.chat(
  { chatPrompt: [{ role: 'user', content: 'Run this when capacity is free.' }] },
  { serviceTier: 'flex' },
);

Inspect ai.getFeatures(model).serviceTiers for the verified tiers on a named profile or exact model. Named OpenAI, Gemini, Azure, Bedrock, Mistral, Groq, OpenRouter, DeepInfra, Cerebras, Grok, Databricks, and Fireworks profiles map the portable value to their provider dialect. Exact modelInfo.supported metadata can opt a conservative custom deployment into a tier.

The applied provider value is normalized into response.modelUsage.tokens.serviceTier; aliases such as default, on_demand, and performance become standard or priority. Add serviceTierPricing.flex or .priority to AxModelInfo when tier pricing differs, and cost estimates will use the tier that actually served the request. Gemini service tiers remain GenerateContent-only: Vertex AI and Gemini Live reject explicit tiers. Anthropic speed: 'fast' remains a separate API.

Routing And Balancing

Choose the primitive by responsibility:

• AxMultiServiceRouter combines model lists and dispatches the model key the caller already selected. It does not select a model.
• AxProviderRouter selects a provider by request capability and may degrade unsupported media through configured processors. When the selected provider supports images natively, each image remains an image object with its payload, MIME type, detail level, cache and optimization hints, alt text, and ordering with surrounding text intact.
• AxBalancer without a strategy orders equivalent services once with a comparator, retries transient provider failures, and fails over in that order.
• AxBalancer with strategy.type: 'adaptive' selects among services exposing the same logical model aliases using learned provider reliability, successful latency, and estimated cost.

Adaptive balancing is operational routing, not semantic prompt-to-model routing. Every provider model behind an alias must be an acceptable substitute for that application. Keep quality evaluation and content-aware model selection outside the balancer.

import { AxBalancer, AxInMemoryBalancerStatsStore } from '@ax-llm/ax';

const statsStore = new AxInMemoryBalancerStatsStore();
const routeKeys = new Map<string, string>([
  [openai.getId(), 'openai-primary'],
  [anthropic.getId(), 'anthropic-primary'],
]);

const llm = AxBalancer.create([openai, anthropic] as const, {
  strategy: {
    type: 'adaptive',
    deadlineMs: 6_000,
    badOutcomeCost: 0.02,
    expectedTokens: { promptTokens: 1_200, completionTokens: 300 },
    namespace: 'support-v1',
    routeKey: (service) => {
      const key = routeKeys.get(service.getId());
      if (!key) throw new Error('Missing stable route key.');
      return key;
    },
    slice: ({ options }) =>
      options?.customLabels?.workflow ?? 'default-workflow',
    statsStore,
    onRoutingEvent: (event) => telemetry.emit('llm.route', event),
  },
});

The score is estimated request cost plus badOutcomeCost times the probability of provider failure or missing deadlineMs. badOutcomeCost and estimated cost must use the same currency or unit. By default, cost uses expectedTokens, the route's concrete model mapping, and getEstimatedCost(); missing catalog pricing contributes zero, while estimateCost can supply application pricing. Failures use an EWMA; successful latency is modeled in log space with a Normal-Inverse-Gamma posterior, and Thompson sampling supplies the deadline risk. Capability filtering still runs before ranking.

Rules:

• Reuse the in-memory store across balancers in one process. For multiple processes, implement AxBalancerStatsStore with Redis or an application database; its observe() operation must be atomic.
• A custom store requires stable, unique routeKey values. Stats are partitioned by namespace, slice, logical model, and route.
• statsStore is decision state. onRoutingEvent is best-effort telemetry and must not be used as the authoritative routing state.
• Routing events contain scores, route metadata, and sanitized failure categories, never prompts, responses, or raw provider errors.
• Adaptive mode attempts each candidate once. Provider-client retries remain controlled by AxAIServiceOptions.retry.
• Streaming can fail over before the first emitted chunk. A mid-stream transient failure is learned but cannot be replayed after partial output; caller cancellation is not recorded.
• Adaptive selection applies to chat. Embedding, transcription, and speech keep existing balancer behavior.

See the adaptive balancer example for complete provider setup.

Chat

const res = await llm.chat({
  chatPrompt: [
    { role: 'system', content: 'You are concise.' },
    { role: 'user', content: 'Write a haiku about the ocean.' },
  ],
});
console.log(res.results[0]?.content);

Batch Audio

Use ai.transcribe(...) for batch speech-to-text and ai.speak(...) for batch text-to-speech. These are separate from conversational .chat() audio config.

const transcript = await llm.transcribe({
  audio: { data: base64Wav, format: 'wav' },
  model: 'gpt-4o-mini-transcribe',
  language: 'en',
});

const speech = await llm.speak({
  text: transcript.text,
  model: 'gpt-4o-mini-tts',
  voice: 'alloy',
  format: 'mp3',
});

console.log(transcript.text);
console.log(speech.data);

Providers without the requested audio endpoint throw AxMediaNotSupportedError. Use speech forward options for signature audio artifacts and modelConfig.audio for conversational chat audio.

Common Options

• stream (boolean): enable SSE; true by default
• thinkingTokenBudget: 'minimal' | 'low' | 'medium' | 'high' | 'highest' | 'none'
• showThoughts: include thoughts in output
• functionCallMode: 'auto' | 'native' | 'prompt'
• debug, logger, tracer, rateLimiter, timeout

Global Runtime Defaults

Use axGlobals when the app wants one live default for AI requests, generator runs, flows, or metrics:

import { ai, axGlobals, axCreateDefaultColorLogger } from '@ax-llm/ax';
import { metrics, trace } from '@opentelemetry/api';

axGlobals.rateLimiter = async (next, info) => next();
axGlobals.tracer = trace.getTracer('my-app');
axGlobals.meter = metrics.getMeter('my-app');
axGlobals.debug = true;
axGlobals.logger = axCreateDefaultColorLogger();
axGlobals.customLabels = { service: 'api' };
axGlobals.onUsage = (event) => usageQueue.enqueue(event);

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

Rules:

• axGlobals.rateLimiter, tracer, meter, logger, debug, abortSignal, and customLabels are live runtime defaults; each operation snapshots them at its start, even if the AI instance already exists.
• Precedence is: per-call options, then explicit AI/service options, then current axGlobals, then built-in defaults.
• The limiter receives next plus operation, provider, model, streaming state, and previous service usage. It wraps chat and embedding provider execution, including streaming and retries; its errors propagate. It may delay, reject, skip, or invoke next multiple times.
• Tracer, meter, and usage-observer failures are fail-open. Limiter failures are fail-closed.
• Runtime-hook telemetry contains metadata and usage only, never prompts, outputs, tool arguments, or tool results.
• External meter instruments are independent of balancer-local getMetrics() snapshots. Adapt AxMeter to OpenTelemetry at the application boundary; generated packages do not require an OpenTelemetry dependency.
• customLabels merge from globals to service to call options; later sources override earlier keys.
• abortSignal values are merged, so either a global shutdown signal or a local request signal can cancel the request.
• axGlobals.onUsage receives one immutable normalized event for each completed chat or embedding call that reports token usage. A fully consumed stream emits once.
• Usage observers are best-effort and fail-open. Ax does not await them; synchronously enqueue events and persist or aggregate them out of band.

Clear process-wide hooks during shutdown or test teardown:

axGlobals.rateLimiter = undefined;
axGlobals.tracer = undefined;
axGlobals.meter = undefined;

Use usageContext for multi-tenant and request attribution:

const llm = ai({
  name: 'openai',
  apiKey: process.env.OPENAI_APIKEY!,
  options: {
    usageContext: {
      tenantId: 'tenant-42',
      feature: 'support-chat',
      attributes: { environment: 'production' },
    },
  },
});

await llm.chat(request, {
  usageContext: {
    userId: user.id,
    requestId: requestId,
    runId: runId,
  },
});

Per-call context overrides service defaults, while attributes are shallow-merged. Events include normalized tokens, provider/model, available session and remote IDs, and a streaming flag. They do not estimate currency cost; calculate that downstream against a versioned pricing table.

DeepSeek Notes

import { ai, AxAIDeepSeekModel } from '@ax-llm/ax';

const deepseek = ai({
  name: 'deepseek',
  apiKey: process.env.DEEPSEEK_APIKEY!,
  config: { model: AxAIDeepSeekModel.DeepSeekV4Flash },
});

DeepSeek's current API models are deepseek-v4-flash and deepseek-v4-pro. Legacy model enum/catalog values remain available for source compatibility, but profile rules are applied only to model IDs verified for the selected deployment.

DeepSeek V4 supports thinking mode. When thinkingTokenBudget is omitted, Ax selects its logical max level and sends thinking: { type: "enabled" } with reasoning_effort: "max". Set thinkingTokenBudget: "none" explicitly to disable it. DeepSeek's API exposes low, medium, high, and max: Ax maps minimal and low to low, preserves medium, maps high to high, and maps highest to max. DeepSeek V4 thinking models support tools, but reject the tool_choice request parameter, so Ax omits auto and Ax-generated __axOutput tool choices for deepseek-v4-pro, deepseek-v4-flash, and deepseek-reasoner while still sending tool definitions. An explicitly forced caller tool choice fails before the request. DeepSeek does not support native JSON schema structured outputs. Ax therefore uses validated json_object for a single required string or code output, including an AxAgent actor's javascriptCode field, without exposing a provider tool. Richer structured outputs use the synthetic __axOutput function when function calling is available, or validated json_object when it is not.

DeepSeek Chat returns thinking traces as reasoning_content. During a tool loop, Ax preserves that field on the assistant tool-call message and sends it back on the following request together with non-null content and the original tool calls. This compatibility mode is declared by the DeepSeek deployment profile; the official OpenAI Chat adapter does not emit or expose reasoning_content.

The same logical default is declared independently for verified DeepSeek V4 rules in the Together, Fireworks, and OpenRouter profiles, then mapped to each deployment's own request dialect. A custom openai-compatible endpoint never inherits it from a DeepSeek-looking model ID.

Other verified deployment rules follow the same policy. Grok 4.6 maps logical max to xhigh, while Grok 4.5 and 4.3 map it to high. Groq GPT-OSS and Cerebras GPT-OSS map it to high; Groq Qwen 3.6 maps it to its documented reasoning-enabled default; Cerebras Gemma 4 and DeepInfra DeepSeek R1 map it to high. An explicit none is sent only for model/deployment combinations that document disabling reasoning. Grok 4.6/4.5 and GPT-OSS on Groq or Cerebras reject none before network I/O because those APIs do not support disabling reasoning for those models.

Hugging Face Router remains conservative: routing policies such as :fastest may choose a different inference provider without changing the base model ID, so Ax does not attach one provider's reasoning contract to that dynamic route.