第 12 章 · 中间件与测试
本章目标:
- 理解 Language Model Middleware 的定位:以模型无关的方式拦截并增强语言模型调用
- 掌握
wrapLanguageModel的用法与多个中间件的执行顺序- 熟悉六个内置中间件及三个自定义实现入口(transformParams / wrapGenerate / wrapStream)
- 学会用
ai/test中的 Mock 模型与测试助手做确定性单元测试- 掌握
simulateReadableStream模拟流式响应的技巧
12.1 什么是 Language Model Middleware
Language model middleware 是一种增强语言模型行为的方式——通过拦截和修改对语言模型的调用来实现。它可以为模型添加 guardrails(护栏)、RAG、缓存、日志等能力,且与具体语言模型无关。这样的中间件可以独立于所应用的模型进行开发和分发。
本章示例中的 yourModel 既可以是 Vercel AI Gateway 构造的模型,也可以是自定义 OpenAI 兼容 Provider 实例——中间件工作在模型抽象层之上,两种接入方式完全通用:
import { createGateway } from 'ai';
const gateway = createGateway({ apiKey: process.env.AI_GATEWAY_API_KEY ?? '' });
const yourModel = gateway('openai/gpt-5');
// 方式二:自定义 OpenAI 兼容 Provider(二选一即可)
// import { createOpenAICompatible } from '@ai-sdk/openai-compatible';
// const yourModel = createOpenAICompatible({
// name: 'my-provider',
// baseURL: process.env.OPENAI_COMPATIBLE_BASE_URL ?? '',
// apiKey: process.env.OPENAI_COMPATIBLE_API_KEY ?? '',
// })('gpt-4o-mini');12.2 使用 wrapLanguageModel
wrapLanguageModel 接收一个语言模型和一个中间件,返回一个融合了该中间件的新模型:
import { wrapLanguageModel, streamText } from 'ai';
const wrappedLanguageModel = wrapLanguageModel({
model: yourModel,
middleware: yourLanguageModelMiddleware,
});包装后的模型可以像普通模型一样使用:
const result = streamText({
model: wrappedLanguageModel,
prompt: 'What cities are in the United States?',
});多个中间件
可以给 wrapLanguageModel 提供多个中间件,按提供的顺序依次应用:
const wrappedLanguageModel = wrapLanguageModel({
model: yourModel,
middleware: [firstMiddleware, secondMiddleware],
});
// 应用顺序: firstMiddleware(secondMiddleware(yourModel))12.3 内置中间件
AI SDK 自带若干内置中间件:
| 中间件 | 用途 |
|---|---|
extractReasoningMiddleware | 从生成文本中提取推理信息(如 <think> 标签),暴露为结果的 reasoning 属性 |
extractJsonMiddleware | 剥离 markdown 代码围栏提取 JSON,配合 Output.object() 使用 |
simulateStreamingMiddleware | 为不支持流式的模型模拟流式行为 |
defaultInstructionsMiddleware | 调用未提供系统指令时应用默认指令 |
defaultSettingsMiddleware | 为模型应用默认设置 |
addToolInputExamplesMiddleware | 把工具输入示例序列化进工具描述 |
提取推理内容
某些 provider 和模型用特殊标签(如 <think> 与 </think>)在生成文本中暴露推理信息:
import { wrapLanguageModel, extractReasoningMiddleware } from 'ai';
const model = wrapLanguageModel({
model: yourModel,
middleware: extractReasoningMiddleware({ tagName: 'think' }),
});还提供 startWithReasoning 选项:设为 true 时会在生成文本前预置推理标签,适用于不把标签放在回复开头的模型。
提取 JSON
有些模型即使你请求结构化输出,也会把 JSON 包在 markdown 代码围栏里。extractJsonMiddleware 会剥离这些围栏使其兼容 Output.object():
import {
wrapLanguageModel,
extractJsonMiddleware,
Output,
generateText,
} from 'ai';
import { z } from 'zod';
const model = wrapLanguageModel({
model: yourModel,
middleware: extractJsonMiddleware(),
});
const result = await generateText({
model,
output: Output.object({
schema: z.object({
name: z.string(),
ingredients: z.array(z.string()),
}),
}),
prompt: 'Generate a recipe.',
});也可以为格式特殊的模型提供自定义 transform 函数:
const model = wrapLanguageModel({
model: yourModel,
middleware: extractJsonMiddleware({
transform: text => text.replace(/^PREFIX/, '').replace(/SUFFIX$/, ''),
}),
});默认设置与默认指令
import { wrapLanguageModel, defaultSettingsMiddleware } from 'ai';
const model = wrapLanguageModel({
model: yourModel,
middleware: defaultSettingsMiddleware({
settings: {
temperature: 0.5,
maxOutputTokens: 800,
providerOptions: { openai: { store: false } },
},
}),
});defaultInstructionsMiddleware 只对没有系统消息的调用应用默认指令;调用中直接给出的指令优先于默认值:
import { wrapLanguageModel, defaultInstructionsMiddleware } from 'ai';
const model = wrapLanguageModel({
model: yourModel,
middleware: defaultInstructionsMiddleware({
instructions: 'You are a concise technical assistant.',
}),
});12.4 实现自定义中间件
⚠️ 实现中间件属于高级功能,需要理解 language model 规范(
LanguageModelV4)。
可以实现以下三个函数之一来修改模型行为:
transformParams:在参数传给模型之前转换它们(对doGenerate和doStream都生效)wrapGenerate:包装模型的doGenerate方法,可修改参数、调用模型并修改结果wrapStream:包装模型的doStream方法,可对流片段做转换
日志中间件
import type {
LanguageModelV4Middleware,
LanguageModelV4StreamPart,
} from '@ai-sdk/provider';
export const yourLogMiddleware: LanguageModelV4Middleware = {
wrapGenerate: async ({ doGenerate, params }) => {
console.log('doGenerate called');
console.log(`params: ${JSON.stringify(params, null, 2)}`);
const result = await doGenerate();
const generatedText = result.content
.filter(part => part.type === 'text')
.map(part => part.text)
.join('');
console.log('doGenerate finished');
console.log(`generated text: ${generatedText}`);
return result;
},
wrapStream: async ({ doStream, params }) => {
console.log('doStream called');
const { stream, ...rest } = await doStream();
let generatedText = '';
const textBlocks = new Map<string, string>();
const transformStream = new TransformStream<
LanguageModelV4StreamPart,
LanguageModelV4StreamPart
>({
transform(chunk, controller) {
switch (chunk.type) {
case 'text-start': {
textBlocks.set(chunk.id, '');
break;
}
case 'text-delta': {
const existing = textBlocks.get(chunk.id) || '';
textBlocks.set(chunk.id, existing + chunk.delta);
generatedText += chunk.delta;
break;
}
case 'text-end': {
console.log(
`Text block ${chunk.id} completed:`,
textBlocks.get(chunk.id),
);
break;
}
}
controller.enqueue(chunk);
},
flush() {
console.log(`generated text: ${generatedText}`);
},
});
return {
stream: stream.pipeThrough(transformStream),
...rest,
};
},
};缓存与 RAG 中间件思路
缓存中间件在 wrapGenerate 中以 JSON.stringify(params) 作为 cache key 命中则直接返回;RAG 中间件则在 transformParams 中取最后一条用户消息、检索相关来源并注入上下文。guardrails 中间件可在 wrapGenerate 中过滤敏感词(如把 badword 替换为 <REDACTED>)。
按请求传递自定义元数据
可以用 providerOptions 在中间件中收发自定义元数据(如用户 ID、时间戳),便于日志追踪:
import { generateText, wrapLanguageModel, createGateway } from 'ai';
import type { LanguageModelV4Middleware } from '@ai-sdk/provider';
import 'dotenv/config';
const gateway = createGateway({ apiKey: process.env.AI_GATEWAY_API_KEY ?? '' });
export const yourLogMiddleware: LanguageModelV4Middleware = {
wrapGenerate: async ({ doGenerate, params }) => {
console.log('METADATA', params?.providerMetadata?.yourLogMiddleware);
const result = await doGenerate();
return result;
},
};
const { text } = await generateText({
model: wrapLanguageModel({
model: gateway('openai/gpt-5'),
middleware: yourLogMiddleware,
}),
prompt: 'Invent a new holiday and describe its traditions.',
providerOptions: {
yourLogMiddleware: {
hello: 'world',
},
},
});此外社区还有遵循 AI SDK 中间件规范的第三方包(如 @ai-sdk-tool/parser),可为不支持原生 function calling 的自托管模型(Hermes/Qwen/Gemma 格式)补上工具调用能力。
12.5 用 ai/test 做确定性测试
测试语言模型代码很难:模型输出非确定、调用又慢又贵。AI SDK Core 提供了 mock providers 和测试助手,可从 ai/test 导入:
MockEmbeddingModelV4:mock 嵌入模型MockLanguageModelV4:mock 语言模型mockId:递增整数 IDmockValues:每次调用迭代返回数组中的值,耗尽后返回最后一个
配合 simulateReadableStream(从 ai 导入)即可在不真正调用 provider 的情况下,以可重复且确定的方式测试你的代码。
测试 generateText
import { generateText } from 'ai';
import { MockLanguageModelV4 } from 'ai/test';
const result = await generateText({
model: new MockLanguageModelV4({
doGenerate: async () => ({
content: [{ type: 'text', text: `Hello, world!` }],
finishReason: { unified: 'stop', raw: undefined },
usage: {
inputTokens: {
total: 10,
noCache: 10,
cacheRead: undefined,
cacheWrite: undefined,
},
outputTokens: {
total: 20,
text: 20,
reasoning: undefined,
},
},
warnings: [],
}),
}),
prompt: 'Hello, test!',
});测试 streamText
import { streamText, simulateReadableStream } from 'ai';
import { MockLanguageModelV4 } from 'ai/test';
const result = streamText({
model: new MockLanguageModelV4({
doStream: async () => ({
stream: simulateReadableStream({
chunks: [
{ type: 'text-start', id: 'text-1' },
{ type: 'text-delta', id: 'text-1', delta: 'Hello' },
{ type: 'text-delta', id: 'text-1', delta: ', ' },
{ type: 'text-delta', id: 'text-1', delta: 'world!' },
{ type: 'text-end', id: 'text-1' },
{
type: 'finish',
finishReason: { unified: 'stop', raw: undefined },
logprobs: undefined,
usage: {
inputTokens: {
total: 3,
noCache: 3,
cacheRead: undefined,
cacheWrite: undefined,
},
outputTokens: {
total: 10,
text: 10,
reasoning: undefined,
},
},
},
],
}),
}),
}),
prompt: 'Hello, test!',
});Mock 同样支持结构化输出场景——让 doGenerate 返回 JSON 字符串文本并传入 output: Output.object({ schema }),或用 simulateReadableStream 分片流出 JSON 片段测试流式对象生成。还可以直接模拟 UI Message Stream 响应(构造 SSE 格式的 data: {...} chunks 并带上 x-vercel-ai-ui-message-stream: v1 头)用于前端调试。
本章小结
- 中间件以模型无关的方式拦截修改调用:
wrapLanguageModel({ model, middleware }),多个中间件按序嵌套 - 六个内置中间件覆盖推理提取、JSON 提取、模拟流式、默认设置、默认指令与工具示例注入
- 自定义中间件有三个入口:
transformParams/wrapGenerate/wrapStream,可实现缓存、RAG、guardrails、日志 providerOptions可向中间件传递按请求的自定义元数据ai/test提供MockLanguageModelV4等 mock 助手 +simulateReadableStream,实现无需真实 provider 的确定性单元测试
🛠️ 动手实践
- 基于 12.4 的日志中间件扩展一个「慢调用告警」:记录
doGenerate耗时,超过 5 秒时打印 warning。 - 给你的模型叠加两个内置中间件(
defaultSettingsMiddleware+extractReasoningMiddleware),并用多中间件写法验证执行顺序。 - 用
MockLanguageModelV4+simulateReadableStream为一个「拼接流式文本并统计长度」的工具函数编写单元测试,断言最终文本与统计值。