第 11 章 · 嵌入向量与重排序

本章目标：

• 理解 embedding 的概念及其在相似度计算与 RAG 中的作用
• 掌握 embed、embedMany 与 cosineSimilarity 三个核心 API
• 学会配置并行请求数、重试次数与超时等生成设置
• 理解 rerank 重排序与向量相似度搜索的区别及适用场景
• 了解嵌入模型中间件 wrapEmbeddingModel 的用法

11.1 什么是 Embedding

Embedding（嵌入）是把词语、短语或图片表示为高维空间中向量的方法。在这个空间里，语义相近的内容彼此靠近，向量间的距离可以用来衡量相似度。

AI SDK 默认通过 AI Gateway 路由模型调用——你可以直接使用 'openai/text-embedding-3-small' 这样的 gateway 模型字符串，也可以换成自定义 OpenAI 兼容 Provider 实例（见第 3 章），两种写法对 embed 系列函数完全透明。

11.2 嵌入单个值

embed 函数用于嵌入单个值，适用于查找相似词句或文本聚类等任务：

import { embed, createGateway } from 'ai';

const gateway = createGateway({ apiKey: process.env.AI_GATEWAY_API_KEY ?? '' });

// 'embedding' 是一个嵌入对象 (number[])
const { embedding } = await embed({
  model: 'openai/text-embedding-3-small', // 经 AI Gateway 路由；也可用自定义 provider 实例
  value: 'sunny day at the beach',
});

11.3 批量嵌入多个值

加载数据时（例如为检索增强生成 RAG 准备数据存储），一次性批量嵌入多个值通常更高效。embedMany 就是为此设计的：

import { embedMany } from 'ai';

// 'embeddings' 是嵌入数组 (number[][])，
// 顺序与输入值一一对应
const { embeddings } = await embedMany({
  model: 'openai/text-embedding-3-small',
  values: [
    'sunny day at the beach',
    'rainy afternoon in the city',
    'snowy night in the mountains',
  ],
});

嵌入相似度

嵌入之后，可以用 cosineSimilarity 函数计算它们之间的余弦相似度，进而对相关条目进行排序和过滤：

import { cosineSimilarity, embedMany } from 'ai';

const { embeddings } = await embedMany({
  model: 'openai/text-embedding-3-small',
  values: ['sunny day at the beach', 'rainy afternoon in the city'],
});

console.log(
  `cosine similarity: ${cosineSimilarity(embeddings[0], embeddings[1])}`,
);

11.4 Token 用量与响应信息

很多 provider 按 token 数量计费。embed 和 embedMany 都会在结果对象的 usage 属性中返回用量信息：

import { embed } from 'ai';

const { embedding, usage } = await embed({
  model: 'openai/text-embedding-3-small',
  value: 'sunny day at the beach',
});

console.log(usage); // { tokens: 10 }

两者还返回包含原始 provider 响应的 response 信息，便于调试：

import { embed } from 'ai';

const { embedding, response } = await embed({
  model: 'openai/text-embedding-3-small',
  value: 'sunny day at the beach',
});

console.log(response); // 原始 provider 响应

11.5 生成设置

Provider Options：通过 providerOptions 配置 provider 专属参数：

import { embed } from 'ai';

const { embedding } = await embed({
  model: 'openai/text-embedding-3-small',
  value: 'sunny day at the beach',
  providerOptions: {
    openai: {
      dimensions: 512, // 降低嵌入维度
    },
  },
});

Google 的 gemini-embedding-2 还支持通过 providerOptions.google.content 进行多模态嵌入，每个条目可含 { text }、{ inlineData } 或 { fileData } 部分。

并行请求：embedMany 支持用 maxParallelCalls 控制并行度以优化性能：

import { embedMany } from 'ai';

const { embeddings, usage } = await embedMany({
  maxParallelCalls: 2, // 限制并行请求数
  model: 'openai/text-embedding-3-small',
  values: [
    'sunny day at the beach',
    'rainy afternoon in the city',
    'snowy night in the mountains',
  ],
});

重试：maxRetries 默认为 2（共尝试 3 次），设为 0 可禁用重试。两者还接受可选的 abortSignal 参数（如 AbortSignal.timeout(1000) 一秒后中止）以及 headers 参数添加自定义请求头。

11.6 嵌入模型中间件

可以用 wrapEmbeddingModel 和 EmbeddingModelMiddleware 增强嵌入模型，例如设置默认值。下面示例使用内置的 defaultEmbeddingSettingsMiddleware：

import {
  defaultEmbeddingSettingsMiddleware,
  embed,
  wrapEmbeddingModel,
  createGateway,
} from 'ai';

const gateway = createGateway({ apiKey: process.env.AI_GATEWAY_API_KEY ?? '' });

const embeddingModelWithDefaults = wrapEmbeddingModel({
  model: gateway.embeddingModel('google/gemini-embedding-001'),
  middleware: defaultEmbeddingSettingsMiddleware({
    settings: {
      providerOptions: {
        google: {
          outputDimensionality: 256,
          taskType: 'CLASSIFICATION',
        },
      },
    },
  }),
});

常用嵌入模型参考：OpenAI text-embedding-3-large（3072 维）/ text-embedding-3-small（1536 维）、Google gemini-embedding-001（3072 维）、Mistral mistral-embed（1024 维）、Cohere embed-multilingual-v3.0（1024 维）等。

11.7 重排序（Reranking）

Reranking 通过按查询相关性对一组文档重新排序来提升搜索质量。与基于 embedding 的相似度搜索不同，reranking 模型经过专门训练以理解 query 与 document 之间的关系，通常能给出更准确的相关性分数。

rerank 函数按相关性对文档重新排序：

import { rerank } from 'ai';

const documents = [
  'sunny day at the beach',
  'rainy afternoon in the city',
  'snowy night in the mountains',
];

const { ranking } = await rerank({
  model: 'cohere/rerank-v3.5', // 经 AI Gateway 路由的重排模型字符串
  documents,
  query: 'talk about rain',
  topN: 2, // 只返回最相关的 2 个文档
});

console.log(ranking);
// [
//   { originalIndex: 1, score: 0.9, document: 'rainy afternoon in the city' },
//   { originalIndex: 0, score: 0.3, document: 'sunny day at the beach' }
// ]

💡 reranking 模型目前主要由 Cohere、Amazon Bedrock、Together.ai 等提供（如 cohere/rerank-v3.5、amazon/cohere.rerank-v3-5:0）。若使用自定义 OpenAI 兼容 Provider，需要服务端实现相应的 reranking 接口。

ranking 数组的每一项包含：originalIndex（原数组中的位置）、score（相关性分数，通常 0–1，越高越相关）、document（原始文档）。结果对象还提供便捷属性：rerankedDocuments（按相关性排序后的文档）与 originalDocuments（原始文档数组）。

对象文档重排序

rerank 也支持结构化文档（JSON 对象），非常适合搜索数据库记录、邮件等内容：

import { rerank } from 'ai';

const documents = [
  {
    from: 'Paul Doe',
    subject: 'Follow-up',
    text: 'We are happy to give you a discount of 20% on your next order.',
  },
  {
    from: 'John McGill',
    subject: 'Missing Info',
    text: 'Sorry, but here is the pricing information from Oracle: $5000/month',
  },
];

const { ranking, rerankedDocuments } = await rerank({
  model: 'cohere/rerank-v3.5',
  documents,
  query: 'Which pricing did we get from Oracle?',
  topN: 1,
});

console.log(rerankedDocuments[0]);
// { from: 'John McGill', subject: 'Missing Info', text: '...' }

设置项

rerank 同样支持 providerOptions（如 Cohere 的 maxTokensPerDoc 限制每文档 token 数）、maxRetries（默认 2 次）、abortSignal 超时控制与 headers 自定义请求头，用法与 embed 系列一致。

import { rerank } from 'ai';

const { ranking } = await rerank({
  model: 'cohere/rerank-v3.5',
  documents: ['doc1', 'doc2', 'doc3', 'doc4', 'doc5'],
  query: 'relevant information',
  topN: 3, // 只返回前 3 个最相关文档
  maxRetries: 0, // 禁用重试
});

本章小结

• Embedding 把内容映射为高维向量，距离即语义相似度；embed 单值、embedMany 批量、cosineSimilarity 算相似度
• usage 返回 token 用量，response 返回原始 provider 响应
• maxParallelCalls / maxRetries / abortSignal / headers / providerOptions 是 embed 系列的通用设置
• wrapEmbeddingModel + 内置中间件可为嵌入模型附加默认设置
• Reranking 由专门训练的模型理解 query-document 关系，通常比纯向量相似度更准；topN 控制返回数量，支持对象文档

🧪 随堂测验

点击你认为正确的选项。答错时会展示正确答案与原因解析。

1. embedMany 返回的 embeddings 数组顺序如何？

2. maxRetries 设为 0 意味着什么？

3. rerank 结果中 ranking 数组每一项不包含下列哪个字段？

4. 关于 reranking 与 embedding 相似度搜索的区别，正确的是？

🛠️ 动手实践

1. 用 embedMany 嵌入 5 句中文短句，再用 cosineSimilarity 计算两两相似度矩阵，找出最相似的一对。
2. 实现「先粗筛再精排」的两阶段搜索：先用 embedMany + 余弦相似度从 50 条数据中取前 10 条，再用 rerank 精排出前 3 条。
3. 用 wrapEmbeddingModel + defaultEmbeddingSettingsMiddleware 封装一个带默认 dimensions 设置的嵌入模型，并验证 usage 中的 token 统计。