Local Caching Middleware

In this implementation, you create a JSON file to store responses. When a request is made, you first check if you have already seen this exact request. If you have, you return the cached response immediately (as a one-off generation or chunks of tokens). If not, you trigger the generation, save the response, and return it.

<Note>

Make sure to add the path of your local cache to your .gitignore so you do

not commit it.

</Note>

### How it works

For regular generations, you store and retrieve complete responses. Instead, the streaming implementation captures each token as it arrives, stores the full sequence, and on cache hits uses the SDK's simulateReadableStream utility to recreate the token-by-token streaming experience at a controlled speed (defaults to 10ms between chunks).

This approach gives you the best of both worlds:

- Instant responses for repeated queries

- Preserved streaming behavior for UI development

The middleware handles all transformations needed to make cached responses indistinguishable from fresh ones, including normalizing tool calls and fixing timestamp formats.

### Middleware

import {

type LanguageModelV3Middleware,

type LanguageModelV3StreamPart,

type LanguageModelV3CallOptions,

type LanguageModelV3,

} from '@ai-toolkit/provider';

import 'dotenv/config';

import fs from 'fs';

import path from 'path';

import { wrapLanguageModel, simulateReadableStream } from 'ai-toolkit';

const CACHE_FILE = path.join(process.cwd(), '.cache/ai-cache.json');

export const cached = (model: LanguageModelV3) =>

wrapLanguageModel({

middleware: cacheMiddleware,

model,

});

const ensureCacheFile = () => {

const cacheDir = path.dirname(CACHE_FILE);

if (!fs.existsSync(cacheDir)) {

fs.mkdirSync(cacheDir, { recursive: true });

}

if (!fs.existsSync(CACHE_FILE)) {

fs.writeFileSync(CACHE_FILE, '{}');

}

};

const getCachedResult = (key: string | object) => {

ensureCacheFile();

const cacheKey = typeof key === 'object' ? JSON.stringify(key) : key;

try {

const cacheContent = fs.readFileSync(CACHE_FILE, 'utf-8');

const cache = JSON.parse(cacheContent);

const result = cache[cacheKey];

return result ?? null;

} catch (error) {

console.error('Cache error:', error);

return null;

}

};

const updateCache = (key: string, value: any) => {

ensureCacheFile();

try {

const cache = JSON.parse(fs.readFileSync(CACHE_FILE, 'utf-8'));

const updatedCache = { ...cache, [key]: value };

fs.writeFileSync(CACHE_FILE, JSON.stringify(updatedCache, null, 2));

} catch (error) {

console.error('Failed to update cache:', error);

}

};

const cleanPrompt = (prompt: LanguageModelV3CallOptions['prompt']) => {

return prompt.map(m => {

if (m.role === 'assistant') {

return {

...m,

content: m.content.map(part =>

part.type === 'tool-call' ? { ...part, toolCallId: 'cached' } : part,

),

};

}

if (m.role === 'tool') {

return {

...m,

content: m.content.map(tc => ({

...tc,

toolCallId: 'cached',

result: {},

})),

};

}

return m;

});

};

export const cacheMiddleware: LanguageModelV3Middleware = {

specificationVersion: 'v3',

wrapGenerate: async ({ doGenerate, params, model }) => {

const cacheKey = JSON.stringify({

prompt: cleanPrompt(params.prompt),

_function: 'generate',

model: model.modelId,

});

const cached = getCachedResult(cacheKey);

if (cached && cached !== null) {

return {

...cached,

response: {

...cached.response,

timestamp: cached?.response?.timestamp

? new Date(cached?.response?.timestamp)

: undefined,

},

};

}

const result = await doGenerate();

updateCache(cacheKey, result);

return result;

},

wrapStream: async ({ doStream, params, model }) => {

const cacheKey = JSON.stringify({

prompt: cleanPrompt(params.prompt),

_function: 'stream',

model: model.modelId,

});

const cached = getCachedResult(cacheKey);

if (cached && cached !== null) {

const { chunks, ...rest } = cached;

const formattedChunks = (chunks as LanguageModelV3StreamPart[]).map(p => {

if (p.type === 'response-metadata' && p.timestamp) {

return { ...p, timestamp: new Date(p.timestamp) };

}

return p;

});

return {

stream: simulateReadableStream({

initialDelayInMs: 0,

chunkDelayInMs: 10,

chunks: formattedChunks,

}),

...rest,

};

}

const { stream, ...rest } = await doStream();

const fullResponse: LanguageModelV3StreamPart[] = [];

const transformStream = new TransformStream<

LanguageModelV3StreamPart,

LanguageModelV3StreamPart

>({

transform(chunk, controller) {

fullResponse.push(chunk);

controller.enqueue(chunk);

},

flush() {

updateCache(cacheKey, { chunks: fullResponse, ...rest });

},

});

return {

stream: stream.pipeThrough(transformStream),

...rest,

};

},

};

The middleware can be easily integrated into your existing AI TOOLKIT setup:

import { openai } from '@ai-toolkit/openai';

import { streamText } from 'ai-toolkit';

import 'dotenv/config';

import { cached } from '../middleware/your-cache-middleware';

async function main() {

const result = streamText({

model: cached(openai('gpt-4o')),

maxOutputTokens: 512,

temperature: 0.3,

maxRetries: 5,

prompt: 'Invent a new holiday and describe its traditions.',

});

for await (const textPart of result.textStream) {

process.stdout.write(textPart);

}

console.log();

console.log('Token usage:', await result.usage);

console.log('Finish reason:', await result.finishReason);

}

main().catch(console.error);

When using this caching middleware, keep these points in mind:

1. Development Only - This approach is intended for local development, not production environments

2. Cache Invalidation - You'll need to clear the cache (delete the cache file) when you want fresh responses

3. Multi-Step Flows - When using stopWhen, be aware that caching occurs at the individual language model response level, not across the entire execution flow. This means that while the model's generation is cached, the tool call is not and will run on each generation.