Skip to content

第 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 实例——中间件工作在模型抽象层之上,两种接入方式完全通用:

ts
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 接收一个语言模型和一个中间件,返回一个融合了该中间件的新模型:

ts
import { wrapLanguageModel, streamText } from 'ai';

const wrappedLanguageModel = wrapLanguageModel({
  model: yourModel,
  middleware: yourLanguageModelMiddleware,
});

包装后的模型可以像普通模型一样使用:

ts
const result = streamText({
  model: wrappedLanguageModel,
  prompt: 'What cities are in the United States?',
});

多个中间件

可以给 wrapLanguageModel 提供多个中间件,按提供的顺序依次应用:

ts
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>)在生成文本中暴露推理信息:

ts
import { wrapLanguageModel, extractReasoningMiddleware } from 'ai';

const model = wrapLanguageModel({
  model: yourModel,
  middleware: extractReasoningMiddleware({ tagName: 'think' }),
});

还提供 startWithReasoning 选项:设为 true 时会在生成文本前预置推理标签,适用于不把标签放在回复开头的模型。

提取 JSON

有些模型即使你请求结构化输出,也会把 JSON 包在 markdown 代码围栏里。extractJsonMiddleware 会剥离这些围栏使其兼容 Output.object()

ts
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 函数:

ts
const model = wrapLanguageModel({
  model: yourModel,
  middleware: extractJsonMiddleware({
    transform: text => text.replace(/^PREFIX/, '').replace(/SUFFIX$/, ''),
  }),
});

默认设置与默认指令

ts
import { wrapLanguageModel, defaultSettingsMiddleware } from 'ai';

const model = wrapLanguageModel({
  model: yourModel,
  middleware: defaultSettingsMiddleware({
    settings: {
      temperature: 0.5,
      maxOutputTokens: 800,
      providerOptions: { openai: { store: false } },
    },
  }),
});

defaultInstructionsMiddleware 只对没有系统消息的调用应用默认指令;调用中直接给出的指令优先于默认值:

ts
import { wrapLanguageModel, defaultInstructionsMiddleware } from 'ai';

const model = wrapLanguageModel({
  model: yourModel,
  middleware: defaultInstructionsMiddleware({
    instructions: 'You are a concise technical assistant.',
  }),
});

12.4 实现自定义中间件

⚠️ 实现中间件属于高级功能,需要理解 language model 规范(LanguageModelV4)。

可以实现以下三个函数之一来修改模型行为:

  1. transformParams:在参数传给模型之前转换它们(对 doGeneratedoStream 都生效)
  2. wrapGenerate:包装模型的 doGenerate 方法,可修改参数、调用模型并修改结果
  3. wrapStream:包装模型的 doStream 方法,可对流片段做转换

日志中间件

ts
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、时间戳),便于日志追踪:

ts
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:递增整数 ID
  • mockValues:每次调用迭代返回数组中的值,耗尽后返回最后一个

配合 simulateReadableStream(从 ai 导入)即可在不真正调用 provider 的情况下,以可重复且确定的方式测试你的代码。

测试 generateText

ts
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

ts
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 的确定性单元测试

🛠️ 动手实践

  1. 基于 12.4 的日志中间件扩展一个「慢调用告警」:记录 doGenerate 耗时,超过 5 秒时打印 warning。
  2. 给你的模型叠加两个内置中间件(defaultSettingsMiddleware + extractReasoningMiddleware),并用多中间件写法验证执行顺序。
  3. MockLanguageModelV4 + simulateReadableStream 为一个「拼接流式文本并统计长度」的工具函数编写单元测试,断言最终文本与统计值。