Chapter 11 · Embeddings and Reranking

Chapter goals:
Understand the concept of embeddings and their role in similarity computation and RAG
Master the three core APIs: embed, embedMany, and cosineSimilarity
Learn to configure generation settings like parallel requests, retries, and timeouts
Understand the difference between reranking and vector similarity search and their applicable scenarios
Learn to use the embedding model middleware wrapEmbeddingModel

11.1 What Is an Embedding

An embedding is a method of representing words, phrases, or images as vectors in high-dimensional space. In this space, semantically similar content sits close together, and the distance between vectors can be used to measure similarity.

AI SDK routes model calls through the AI Gateway by default — you can use gateway model strings like 'openai/text-embedding-3-small' directly, or swap in a custom OpenAI-compatible Provider instance (see Chapter 3); both approaches are fully transparent to the embed series functions.

11.2 Embedding a Single Value

The embed function embeds a single value, suitable for tasks like finding similar phrases or text clustering:

import { embed, createGateway } from 'ai';

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

// 'embedding' is an embedding object (number[])
const { embedding } = await embed({
  model: 'openai/text-embedding-3-small', // Routed via AI Gateway; or use a custom provider instance
  value: 'sunny day at the beach',
});

11.3 Batch-Embedding Multiple Values

When loading data (e.g., preparing a data store for Retrieval-Augmented Generation / RAG), batch-embedding multiple values at once is通常 more efficient. embedMany is designed for this:

import { embedMany } from 'ai';

// 'embeddings' is an array of embeddings (number[][]),
// ordered to correspond one-to-one with the input values
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',
  ],
});

Embedding Similarity

After embedding, you can use the cosineSimilarity function to compute the cosine similarity between them, then sort and filter related entries:

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 Usage and Response Info

Many providers bill by token count. Both embed and embedMany return usage information in the usage property of the result object:

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 }

They also return response info containing the raw provider response,便于 debugging:

import { embed } from 'ai';

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

console.log(response); // Raw provider response

11.5 Generation Settings

Provider Options: Configure provider-specific parameters via providerOptions:

import { embed } from 'ai';

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

Google's gemini-embedding-2 also supports multimodal embeddings via providerOptions.google.content, where each entry can contain { text }, { inlineData }, or { fileData } parts.

Parallel Requests: embedMany supports controlling concurrency with maxParallelCalls to optimize performance:

import { embedMany } from 'ai';

const { embeddings, usage } = await embedMany({
  maxParallelCalls: 2, // Limit parallel request count
  model: 'openai/text-embedding-3-small',
  values: [
    'sunny day at the beach',
    'rainy afternoon in the city',
    'snowy night in the mountains',
  ],
});

Retries: maxRetries defaults to 2 (3 total attempts); set to 0 to disable. Both also accept an optional abortSignal parameter (e.g., AbortSignal.timeout(1000) to abort after 1 second) and a headers parameter to add custom request headers.

11.6 Embedding Model Middleware

You can enhance embedding models using wrapEmbeddingModel and EmbeddingModelMiddleware, e.g., to set defaults. The following example uses the built-in 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',
        },
      },
    },
  }),
});

Common embedding models: OpenAI text-embedding-3-large (3072 dims) / text-embedding-3-small (1536 dims), Google gemini-embedding-001 (3072 dims), Mistral mistral-embed (1024 dims), Cohere embed-multilingual-v3.0 (1024 dims), etc.

11.7 Reranking

Reranking improves search quality by reordering a set of documents by query relevance. Unlike embedding-based similarity search, reranking models are专门 trained to understand the relationship between a query and documents,通常 yielding more accurate relevance scores.

The rerank function reranks documents by relevance:

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: 'openai/rerank-2',
  query: 'beach weather',
  documents,
});

The result contains a ranking array with each document's relevance score, ordered from highest to lowest.

Chapter Summary

Embeddings convert text/images into high-dimensional vectors; semantic similarity is measured by vector distance;
Core APIs: embed (single value), embedMany (batch), cosineSimilarity (similarity computation);
Generation settings: providerOptions for provider-specific params, maxParallelCalls for concurrency, maxRetries and abortSignal for resilience;
wrapEmbeddingModel + middleware lets you set defaults and enhance embedding behavior;
Reranking reorders documents by query relevance using a dedicated model — more accurate than pure embedding similarity for retrieval.

🧪 随堂测验

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

1. What does cosineSimilarity return?

2. What is the difference between embed and embedMany?

3. What does the rerank function do that embedding similarity alone does not?

4. How do you reduce the embedding dimension for OpenAI's text-embedding-3-small?

🛠️ Hands-On Exercises

Embed three sentences with embedMany; compute pairwise cosine similarity and print the most similar pair.
Use rerank to rank five documents against a query; compare the reranked order with the order from raw cosine similarity.
Wrap an embedding model with defaultEmbeddingSettingsMiddleware to set a default dimension of 256; verify the setting applies.

Next: Chapter 12 · Middleware and Testing

Chapter 11 · Embeddings and Reranking ​

11.1 What Is an Embedding ​

11.2 Embedding a Single Value ​

11.3 Batch-Embedding Multiple Values ​

Embedding Similarity ​

11.4 Token Usage and Response Info ​

11.5 Generation Settings ​

11.6 Embedding Model Middleware ​

11.7 Reranking ​

Chapter Summary ​