How Context Caching Works with Different LLM Providers in Ax

Ax implements provider-agnostic context caching by allowing developers to mark reusable prompt segments with a .cache() modifier, which each LLM provider then serializes into its native cache-control format.

The ax-llm/ax library provides a unified interface for building LLM applications across multiple providers. Context caching in Ax lets you mark specific parts of your prompts—such as system instructions or large static contexts—as reusable across multiple API calls, significantly reducing token costs and latency while maintaining provider-specific optimizations.

The Five-Step Context Caching Workflow

Ax implements context caching through a layered architecture that abstracts provider differences while preserving native performance benefits.

1. Mark Fields for Caching with .cache()

In src/ax/dsp/sig.ts (around line 583), you use the fluent .cache() modifier on signature fields to indicate reusability. When you define a field like f.string().cache(), Ax emits this field in the request payload with a cache: true flag.

2. Detect Caching in the Base AI Class

The generic AI base class in src/ax/ai/base.ts (lines 131-138) scans every message and tool payload for the cache: true flag. If detected, the request is treated as cache-enabled, and the system prepares to inject appropriate cache-control directives.

3. Provider-Specific Encoding

Each LLM provider receives the caching hint in its native format through specialized adapters:

Anthropic Implementation

In src/ax/ai/anthropic/api.ts (lines 248-1102), the Anthropic-specific implementation converts the cache: true flag into cache_control: {type: 'ephemeral'} objects attached to message or tool blocks. This instructs Anthropic's API to store the content embedding for later reuse within the conversation window.

Google Gemini Implementation

In src/ax/ai/google-gemini/api.ts (lines 1411-1519), the Gemini implementation automatically detects the breakpoint—the last cached part in the sequence—and injects a cache_control field at that position in the request payload. Gemini natively caches system prompts and explicitly marked segments.

Other Providers

For providers that do not expose a cache_control field (such as OpenAI's standard /v1/chat/completions endpoint), Ax still passes the cache: true flag in the metadata. The caching abstraction degrades gracefully, resulting in a standard request without provider-side caching, while Ax's client-side caching remains functional.

4. Cache Key Generation and Storage

When a flow or generator runs, Ax computes a deterministic cache key from the cached values using the getCacheKey function in src/ax/flow/flow.ts (lines 991-1006). The LLM's response is then stored in your configured cache implementation—whether an in-memory Map, Redis, or another provider. Subsequent calls with identical cached content retrieve the result from storage instead of invoking the model.

5. TTL Management and Refresh

For providers supporting ephemeral caches, Ax automatically refreshes the time-to-live (TTL) on every cache hit. This keeps static context alive for the duration of long conversations without requiring manual cache management.

Practical Implementation Examples

Anthropic Context Caching Example

The following example demonstrates caching a large static context with Anthropic's Claude:

import { ai, agent, f, s } from '@ax-llm/ax';

// Create an Anthropic AI instance
const anthropic = ai({
  name: 'anthropic',
  apiKey: process.env.ANTHROPIC_APIKEY!,
});

// Define a signature; mark the rarely-changing context as cacheable
const sig = s(`
  staticContext:string "Static context" .cache()
  userQuestion:string "Current user query"
`);

// Build an agent
const myAgent = agent(sig, {
  name: 'cacheDemo',
  description: 'Demo of context caching with Anthropic',
  ai: anthropic,
});

// Run two turns – the second hit reuses the cached staticContext
await myAgent.run({ 
  staticContext: 'Company policy v1', 
  userQuestion: 'What is the refund policy?' 
});

await myAgent.run({ 
  staticContext: 'Company policy v1', 
  userQuestion: 'Can I get a refund for a cancelled order?' 
});

On the first call, the staticContext field is sent with cache: true, and Anthropic receives the cache_control: {type: 'ephemeral'} directive, storing the embedding. On the second call, Ax detects the same cache key, skips the API request for the cached portion, and only submits the new userQuestion, resulting in reduced token usage and lower latency.

Google Gemini Context Caching Example

Gemini automatically handles system prompt caching and respects explicit breakpoints:

import { ai, agent, f, s } from '@ax-llm/ax';

const gemini = ai({
  name: 'google-gemini',
  apiKey: process.env.GEMINI_APIKEY!,
});

const sig = s(`
  systemPrompt:string "System instructions" .cache()
  userMessage:string "User input"
`);

const geminiAgent = agent(sig, {
  name: 'geminiCacheDemo',
  description: 'Gemini context caching example',
  ai: gemini,
});

await geminiAgent.run({ 
  systemPrompt: 'You are a helpful assistant.', 
  userMessage: 'Explain recursion.' 
});

await geminiAgent.run({ 
  systemPrompt: 'You are a helpful assistant.', 
  userMessage: 'Now give a code example.' 
});

In this implementation, Gemini caches the system prompt automatically. The second turn reuses the cached content, and the request payload includes a cache_control entry for the breakpoint at the last cached message, as implemented in the Gemini-specific provider code.

Summary

• Mark fields for caching using the .cache() fluent modifier in your signatures to designate reusable content.
• Provider-specific serialization occurs automatically in src/ax/ai/anthropic/api.ts and src/ax/ai/google-gemini/api.ts, converting Ax's generic cache: true flag into native cache_control directives.
• Graceful degradation ensures that unsupported providers receive standard requests without breaking the application flow.
• Deterministic cache keys generated in src/ax/flow/flow.ts enable client-side caching with configurable storage backends.
• Automatic TTL refresh maintains ephemeral caches throughout long-running conversations.

Frequently Asked Questions

What is context caching in Ax?

Context caching is a provider-agnostic feature in the Ax framework that allows you to mark specific portions of a prompt—such as system instructions, large documents, or static context—as reusable across multiple LLM calls. According to the ax-llm/ax source code, this reduces API costs and latency by avoiding redundant processing of identical content.

How do I enable context caching for a specific field?

You enable caching by appending the .cache() modifier to any field definition in your signature, such as f.string().cache() or within the signature template syntax. This modifier, implemented in src/ax/dsp/sig.ts, flags the field for provider-specific cache optimization.

Does context caching work with OpenAI models?

Context caching degrades gracefully with OpenAI and other providers that do not support explicit cache-control headers. While OpenAI's standard API ignores the cache: true flag, Ax still maintains client-side caching through the deterministic cache key system in src/ax/flow/flow.ts, preserving some performance benefits.

How does Ax handle cache expiration?

For providers like Anthropic that support ephemeral caching, Ax automatically refreshes the TTL on every cache hit, keeping the cached content alive throughout the conversation. For client-side caches using custom implementations (Redis, in-memory, etc.), expiration is controlled by your configured cache backend settings.

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