feat (ai/core): add embed function (#1575)

Author: lgrammel

.changeset/witty-beds-sell.md

@@ -0,0 +1,8 @@
+---
+'@ai-sdk/provider': patch
+'@ai-sdk/mistral': patch
+'@ai-sdk/openai': patch
+'ai': patch
+---
+
+feat (ai/core): add embed function

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

@@ -0,0 +1,24 @@
+---
+title: Embeddings
+description: Learn how to embed values with the Vercel AI SDK.
+---
+
+# Embeddings
+
+Embeddings are a way to represent words, phrases, or images as vectors in a high-dimensional space.
+In this space, similar words are close to each other, and the distance between words can be used to measure their similarity.
+
+## 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')`.
+
+```tsx
+import { embed } from 'ai';
+import { openai } from '@ai-sdk/openai';
+
+const { embedding } = await embed({
+  model: openai.embedding('text-embedding-3-small'),
+  value: 'sunny day at the beach',
+});
+```

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

@@ -0,0 +1,16 @@
+import { mistral } from '@ai-sdk/mistral';
+import { embed } from 'ai';
+import dotenv from 'dotenv';
+
+dotenv.config();
+
+async function main() {
+  const { embedding } = await embed({
+    model: mistral.embedding('mistral-embed'),
+    value: 'sunny day at the beach',
+  });
+
+  console.log(embedding);
+}
+
+main().catch(console.error);

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

@@ -0,0 +1,16 @@
+import { openai } from '@ai-sdk/openai';
+import { embed } from 'ai';
+import dotenv from 'dotenv';
+
+dotenv.config();
+
+async function main() {
+  const { embedding } = await embed({
+    model: openai.embedding('text-embedding-3-small'),
+    value: 'sunny day at the beach',
+  });
+
+  console.log(embedding);
+}
+
+main().catch(console.error);

packages/core/core/embed/embed.test.ts

@@ -0,0 +1,25 @@
+import assert from 'node:assert';
+import { MockEmbeddingModelV1 } from '../test/mock-embedding-model-v1';
+import { embed } from './embed';
+
+const dummyEmbedding = [0.1, 0.2, 0.3];
+const testValue = 'sunny day at the beach';
+
+describe('result.embedding', () => {
+  it('should generate embedding', async () => {
+    const result = await embed({
+      model: new MockEmbeddingModelV1({
+        doEmbed: async ({ values }) => {
+          assert.deepStrictEqual(values, [testValue]);
+
+          return {
+            embeddings: [dummyEmbedding],
+          };
+        },
+      }),
+      value: testValue,
+    });
+
+    assert.deepStrictEqual(result.embedding, dummyEmbedding);
+  });
+});

packages/core/core/embed/embed.ts

@@ -0,0 +1,95 @@
+import { Embedding, EmbeddingModel } from '../types';
+import { retryWithExponentialBackoff } from '../util/retry-with-exponential-backoff';
+
+/**
+Embed a value using an embedding model. The type of the value is defined by the embedding model.
+
+@param model - The embedding model to use.
+@param value - The value that should be embedded.
+
+@param maxRetries - Maximum number of retries. Set to 0 to disable retries. Default: 2.
+@param abortSignal - An optional abort signal that can be used to cancel the call.
+
+@returns A result object that contains the embedding, the value, and additional information.
+ */
+export async function embed<VALUE>({
+  model,
+  value,
+  maxRetries,
+  abortSignal,
+}: {
+  /**
+The embedding model to use.
+     */
+  model: EmbeddingModel<VALUE>;
+
+  /**
+The value that should be embedded.
+   */
+  value: VALUE;
+
+  /**
+Maximum number of retries per embedding model call. Set to 0 to disable retries.
+
+@default 2
+   */
+  maxRetries?: number;
+
+  /**
+Abort signal.
+ */
+  abortSignal?: AbortSignal;
+}): Promise<EmbedResult<VALUE>> {
+  const retry = retryWithExponentialBackoff({ maxRetries });
+
+  const modelResponse = await retry(() =>
+    model.doEmbed({
+      values: [value],
+      abortSignal,
+    }),
+  );
+
+  return new EmbedResult({
+    value,
+    embedding: modelResponse.embeddings[0],
+    rawResponse: modelResponse.rawResponse,
+  });
+}
+
+/**
+The result of a `embed` call.
+It contains the embedding, the value, and additional information.
+ */
+export class EmbedResult<VALUE> {
+  /**
+The value that was embedded.
+   */
+  readonly value: VALUE;
+
+  /**
+The embedding of the value.
+  */
+  readonly embedding: Embedding;
+
+  /**
+Optional raw response data.
+   */
+  readonly rawResponse?: {
+    /**
+Response headers.
+     */
+    headers?: Record<string, string>;
+  };
+
+  constructor(options: {
+    value: VALUE;
+    embedding: Embedding;
+    rawResponse?: {
+      headers?: Record<string, string>;
+    };
+  }) {
+    this.value = options.value;
+    this.embedding = options.embedding;
+    this.rawResponse = options.rawResponse;
+  }
+}

packages/core/core/embed/index.ts

@@ -0,0 +1 @@
+export * from './embed';

packages/core/core/index.ts

@@ -1,3 +1,4 @@
+export * from './embed';
 export * from './generate-object';
 export * from './generate-text';
 export * from './prompt';

packages/core/core/test/mock-embedding-model-v1.ts

@@ -0,0 +1,36 @@
+import { EmbeddingModelV1 } from '@ai-sdk/provider';
+
+export class MockEmbeddingModelV1<VALUE> implements EmbeddingModelV1<VALUE> {
+  readonly specificationVersion = 'v1';
+
+  readonly provider: EmbeddingModelV1<VALUE>['provider'];
+  readonly modelId: EmbeddingModelV1<VALUE>['modelId'];
+  readonly maxEmbeddingsPerCall: EmbeddingModelV1<VALUE>['maxEmbeddingsPerCall'];
+  readonly supportsParallelCalls: EmbeddingModelV1<VALUE>['supportsParallelCalls'];
+
+  doEmbed: EmbeddingModelV1<VALUE>['doEmbed'];
+
+  constructor({
+    provider = 'mock-provider',
+    modelId = 'mock-model-id',
+    maxEmbeddingsPerCall = 1,
+    supportsParallelCalls = false,
+    doEmbed = notImplemented,
+  }: {
+    provider?: EmbeddingModelV1<VALUE>['provider'];
+    modelId?: EmbeddingModelV1<VALUE>['modelId'];
+    maxEmbeddingsPerCall?: EmbeddingModelV1<VALUE>['maxEmbeddingsPerCall'];
+    supportsParallelCalls?: EmbeddingModelV1<VALUE>['supportsParallelCalls'];
+    doEmbed?: EmbeddingModelV1<VALUE>['doEmbed'];
+  }) {
+    this.provider = provider;
+    this.modelId = modelId;
+    this.maxEmbeddingsPerCall = maxEmbeddingsPerCall;
+    this.supportsParallelCalls = supportsParallelCalls;
+    this.doEmbed = doEmbed;
+  }
+}
+
+function notImplemented(): never {
+  throw new Error('Not implemented');
+}

packages/core/core/types/embedding-model.ts

@@ -0,0 +1,11 @@
+import { EmbeddingModelV1, EmbeddingModelV1Embedding } from '@ai-sdk/provider';
+
+/**
+Embedding model that is used by the AI SDK Core functions.
+*/
+export type EmbeddingModel<VALUE> = EmbeddingModelV1<VALUE>;
+
+/**
+Embedding.
+ */
+export type Embedding = EmbeddingModelV1Embedding;