feat (ai/core): add embedMany function (#1617)

.changeset/five-knives-deny.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+feat (ai/core): add embedMany function

content/docs/03-ai-sdk-core/30-embeddings.mdx

@@ -10,15 +10,42 @@ In this space, similar words are close to each other, and the distance between w
 
 ## Embedding a Single Value
 
-The Vercel AI SDK provides the `embed` function to embed single values, which is useful for tasks such as finding similar words
-or phrases or clustering text. You can use it with embeddings models, e.g. `openai.embedding('text-embedding-3-large')` or `mistral.embedding('mistral-embed')`.
+The Vercel AI SDK provides the [`embed`](/docs/reference/ai-sdk-core/embed) function to embed single values, which is useful for tasks such as finding similar words
+or phrases or clustering text.
+You can use it with embeddings models, e.g. `openai.embedding('text-embedding-3-large')` or `mistral.embedding('mistral-embed')`.
 
 ```tsx
 import { embed } from 'ai';
 import { openai } from '@ai-sdk/openai';
 
+// 'embedding' is a single embedding object (number[])
 const { embedding } = await embed({
   model: openai.embedding('text-embedding-3-small'),
   value: 'sunny day at the beach',
 });
 ```
+
+## Embedding Many Values
+
+When loading data, e.g. when preparing a data store for retrieval-augmented generation (RAG),
+it is often useful to embed many values at once (batch embedding).
+
+The Vercel AI SDK provides the `embedMany` function for this purpose.
+Similar to `embed`, you can use it with embeddings models,
+e.g. `openai.embedding('text-embedding-3-large')` or `mistral.embedding('mistral-embed')`.
+
+```tsx
+import { openai } from '@ai-sdk/openai';
+import { embedMany } from 'ai';
+
+// 'embeddings' is an array of embedding objects (number[][]).
+// It is sorted in the same order as the input values.
+const { embeddings } = await embedMany({
+  model: openai.embedding('text-embedding-3-small'),
+  values: [
+    'sunny day at the beach',
+    'rainy afternoon in the city',
+    'snowy night in the mountains',
+  ],
+});
+```

content/docs/07-reference/ai-sdk-core/06-embed-many.mdx

@@ -0,0 +1,18 @@
+---
+title: embedMany
+description: Embed several values using the AI SDK Core (batch embedding)
+---
+
+# `embedMany`
+
+Embed several values using an embedding model. The type of the value is defined
+by the embedding model.
+
+`embedMany` automatically splits large requests into smaller chunks if the model
+has a limit on how many embeddings can be generated in a single call.
+
+## Import
+
+<Snippet text={`import { embedMany } from "ai"`} prompt={false} />
+
+<ReferenceTable packageName="core" functionName="embedMany" />

content/docs/07-reference/ai-sdk-core/index.mdx

@@ -33,5 +33,11 @@ description: Reference documentation for the AI SDK Core
         'Generate an embedding for a single value using an embedding model.',
       href: '/docs/reference/ai-sdk-core/embed',
     },
+    {
+      title: 'embedMany',
+      description:
+        'Generate embeddings for several values using an embedding model (batch embedding).',
+      href: '/docs/reference/ai-sdk-core/embed-many',
+    },
   ]}
 />

content/providers/01-ai-sdk-providers/01-openai.mdx

@@ -72,9 +72,9 @@ You can use the following optional settings to customize the OpenAI provider ins
   and `compatible` when using 3rd party providers. In `compatible` mode, newer
   information such as streamOptions are not being sent. Defaults to 'compatible'.
 
-## Models
+## Language Models
 
-The OpenAI provider instance is a function that you can invoke to create a model:
+The OpenAI provider instance is a function that you can invoke to create a language model:
 
 ```ts
 const model = openai('gpt-3.5-turbo');
@@ -92,6 +92,14 @@ const model = openai('gpt-3.5-turbo', {
 The available options depend on the API that's automatically chosen for the model (see below).
 If you want to explicitly select a specific model API, you can use `.chat` or `.completion`.
 
+### Model Capabilities
+
+| Model           | Image Input         | Object Generation   | Tool Usage          | Tool Streaming      |
+| --------------- | ------------------- | ------------------- | ------------------- | ------------------- |
+| `gpt-4-turbo`   | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
+| `gpt-4`         | <Cross size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
+| `gpt-3.5-turbo` | <Cross size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
+
 ### Chat Models
 
 You can create models that call the [OpenAI chat API](https://platform.openai.com/docs/api-reference/chat) using the `.chat()` factory method.
@@ -215,10 +223,40 @@ The following optional settings are available for OpenAI completion models:
   A unique identifier representing your end-user, which can help OpenAI to
   monitor and detect abuse. Learn more.
 
-## Model Capabilities
+## Embedding Models
 
-| Model           | Image Input         | Object Generation   | Tool Usage          | Tool Streaming      |
-| --------------- | ------------------- | ------------------- | ------------------- | ------------------- |
-| `gpt-4-turbo`   | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
-| `gpt-4`         | <Cross size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
-| `gpt-3.5-turbo` | <Cross size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
+You can create models that call the [OpenAI embeddings API](https://platform.openai.com/docs/api-reference/embeddings)
+using the `.embedding()` factory method.
+
+```ts
+const model = openai.embedding('text-embedding-3-large');
+```
+
+OpenAI embedding models support several aditional settings.
+You can pass them as an options argument:
+
+```ts
+const model = openai.embedding('text-embedding-3-large', {
+  dimensions: 512 // optional, number of dimensions for the embedding
+  user: 'test-user' // optional unique user identifier
+})
+```
+
+The following optional settings are available for OpenAI embedding models:
+
+- **dimensions**: _number_
+
+  Echo back the prompt in addition to the completion.
+
+- **user** _string_
+
+  A unique identifier representing your end-user, which can help OpenAI to
+  monitor and detect abuse. Learn more.
+
+### Model Capabilities
+
+| Model                    | Default Dimensions | Custom Dimensions   |
+| ------------------------ | ------------------ | ------------------- |
+| `text-embedding-3-large` | 3072               | <Check size={18} /> |
+| `text-embedding-3-small` | 1536               | <Check size={18} /> |
+| `text-embedding-ada-002` | 1536               | <Cross size={18} /> |

content/providers/01-ai-sdk-providers/02-anthropic.mdx

@@ -60,7 +60,7 @@ You can use the following optional settings to customize the Google Generative A
 
   Custom headers to include in the requests.
 
-## Models
+## Language Models
 
 You can create models that call the [Anthropic Messages API](https://docs.anthropic.com/claude/reference/messages_post) using the provider instance.
 The first argument is the model id, e.g. `claude-3-haiku-20240307`.
@@ -88,7 +88,7 @@ The following optional settings are available for Anthropic models:
   Used to remove "long tail" low probability responses.
   Recommended for advanced use cases only. You usually only need to use temperature.
 
-## Model Capabilities
+### Model Capabilities
 
 | Model                      | Image Input         | Object Generation   | Tool Usage          | Tool Streaming      |
 | -------------------------- | ------------------- | ------------------- | ------------------- | ------------------- |

content/providers/01-ai-sdk-providers/03-google-generative-ai.mdx

@@ -58,7 +58,7 @@ You can use the following optional settings to customize the Google Generative A
 
   Custom headers to include in the requests.
 
-## Models
+## Language Models
 
 You can create models that call the [Google Generative AI API](https://ai.google.dev/api/rest) using the provider instance.
 The first argument is the model id, e.g. `models/gemini-pro`.
@@ -87,7 +87,7 @@ The following optional settings are available for Google Generative AI models:
   Top-k sampling considers the set of topK most probable tokens.
   Models running with nucleus sampling don't allow topK setting.
 
-## Model Capabilities
+### Model Capabilities
 
 | Model                          | Image Input         | Object Generation   | Tool Usage          | Tool Streaming      |
 | ------------------------------ | ------------------- | ------------------- | ------------------- | ------------------- |

content/providers/01-ai-sdk-providers/04-mistral.mdx

@@ -23,6 +23,7 @@ The Mistral provider is available in the `@ai-sdk/mistral` module. You can insta
     <Snippet text="yarn add @ai-sdk/mistral" dark />
   </Tab>
 </Tabs>
+
 ## Provider Instance
 
 You can import the default provider instance `mistral` from `@ai-sdk/mistral`:
@@ -58,7 +59,7 @@ You can use the following optional settings to customize the Mistral provider in
 
   Custom headers to include in the requests.
 
-## Models
+## Language Models
 
 You can create models that call the [Mistral chat API](https://docs.mistral.ai/api/#operation/createChatCompletion) using provider instance.
 The first argument is the model id, e.g. `mistral-large-latest`.
@@ -85,9 +86,18 @@ The following optional settings are available for Mistral models:
 
   Defaults to `false`.
 
-## Model Capabilities
+### Model Capabilities
 
 | Model                  | Image Input         | Object Generation   | Tool Usage          | Tool Streaming      |
 | ---------------------- | ------------------- | ------------------- | ------------------- | ------------------- |
 | `mistral-large-latest` | <Cross size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
 | `mistral-small-latest` | <Cross size={18} /> | <Check size={18} /> | <Check size={18} /> | <Check size={18} /> |
+
+## Embedding Models
+
+You can create models that call the [Mistral embeddings API](https://docs.mistral.ai/api/#operation/createEmbedding)
+using the `.embedding()` factory method.
+
+```ts
+const model = mistral.embedding('mistral-embed');
+```

examples/ai-core/src/embed-many/mistral.ts

@@ -0,0 +1,20 @@
+import { mistral } from '@ai-sdk/mistral';
+import { embedMany } from 'ai';
+import dotenv from 'dotenv';
+
+dotenv.config();
+
+async function main() {
+  const { embeddings } = await embedMany({
+    model: mistral.embedding('mistral-embed'),
+    values: [
+      'sunny day at the beach',
+      'rainy afternoon in the city',
+      'snowy night in the mountains',
+    ],
+  });
+
+  console.log(embeddings);
+}
+
+main().catch(console.error);

examples/ai-core/src/embed-many/openai.ts

@@ -0,0 +1,20 @@
+import { openai } from '@ai-sdk/openai';
+import { embedMany } from 'ai';
+import dotenv from 'dotenv';
+
+dotenv.config();
+
+async function main() {
+  const { embeddings } = await embedMany({
+    model: openai.embedding('text-embedding-3-small'),
+    values: [
+      'sunny day at the beach',
+      'rainy afternoon in the city',
+      'snowy night in the mountains',
+    ],
+  });
+
+  console.log(embeddings);
+}
+
+main().catch(console.error);

packages/core/core/embed/embed-many.test.ts

@@ -0,0 +1,74 @@
+import assert from 'node:assert';
+import {
+  MockEmbeddingModelV1,
+  mockEmbed,
+} from '../test/mock-embedding-model-v1';
+import { embedMany } from './embed-many';
+
+const dummyEmbeddings = [
+  [0.1, 0.2, 0.3],
+  [0.4, 0.5, 0.6],
+  [0.7, 0.8, 0.9],
+];
+
+const testValues = [
+  'sunny day at the beach',
+  'rainy afternoon in the city',
+  'snowy night in the mountains',
+];
+
+describe('result.embedding', () => {
+  it('should generate embeddings', async () => {
+    const result = await embedMany({
+      model: new MockEmbeddingModelV1({
+        maxEmbeddingsPerCall: 5,
+        doEmbed: mockEmbed(testValues, dummyEmbeddings),
+      }),
+      values: testValues,
+    });
+
+    assert.deepStrictEqual(result.embeddings, dummyEmbeddings);
+  });
+
+  it('should generate embeddings when several calls are required', async () => {
+    let callCount = 0;
+
+    const result = await embedMany({
+      model: new MockEmbeddingModelV1({
+        maxEmbeddingsPerCall: 2,
+        doEmbed: async ({ values }) => {
+          if (callCount === 0) {
+            assert.deepStrictEqual(values, testValues.slice(0, 2));
+            callCount++;
+            return { embeddings: dummyEmbeddings.slice(0, 2) };
+          }
+
+          if (callCount === 1) {
+            assert.deepStrictEqual(values, testValues.slice(2));
+            callCount++;
+            return { embeddings: dummyEmbeddings.slice(2) };
+          }
+
+          throw new Error('Unexpected call');
+        },
+      }),
+      values: testValues,
+    });
+
+    assert.deepStrictEqual(result.embeddings, dummyEmbeddings);
+  });
+});
+
+describe('result.values', () => {
+  it('should include values in the result', async () => {
+    const result = await embedMany({
+      model: new MockEmbeddingModelV1({
+        maxEmbeddingsPerCall: 5,
+        doEmbed: mockEmbed(testValues, dummyEmbeddings),
+      }),
+      values: testValues,
+    });
+
+    assert.deepStrictEqual(result.values, testValues);
+  });
+});