feat (core): add text embedding model support to provider registry (#…

…1959)

Co-authored-by: Grace Yun <74513600+iteratetograceness@users.noreply.github.com>

.changeset/brave-poets-worry.md

@@ -0,0 +1,11 @@
+---
+'@ai-sdk/google-vertex': patch
+'@ai-sdk/anthropic': patch
+'@ai-sdk/mistral': patch
+'@ai-sdk/google': patch
+'@ai-sdk/openai': patch
+'@ai-sdk/azure': patch
+'ai': patch
+---
+
+feat (core): add text embedding model support to provider registry

content/docs/03-ai-sdk-core/40-provider-management.mdx

@@ -14,7 +14,7 @@ The Vercel AI SDK provides a [`ProviderRegistry`](/docs/reference/ai-sdk-core/pr
 You can register multiple providers. The provider id will become the prefix of the model id:
 `providerId:modelId`.
 
-### Setup (Example)
+### Setup
 
 You can create a registry with multiple providers and models using `experimental_createProviderRegistry`.
 
@@ -39,7 +39,7 @@ export const registry = createProviderRegistry({
 });
 ```
 
-### Usage (Example)
+### Language models
 
 You can access language models by using the `languageModel` method on the registry.
 The provider id will become the prefix of the model id: `providerId:modelId`.
@@ -53,3 +53,18 @@ const { text } = await generateText({
   prompt: 'Invent a new holiday and describe its traditions.',
 });
 ```
+
+### Text embedding models
+
+You can access text embedding models by using the `textEmbeddingModel` method on the registry.
+The provider id will become the prefix of the model id: `providerId:modelId`.
+
+```ts highlight={"5"}
+import { embed } from 'ai';
+import { registry } from './registry';
+
+const { embedding } = await embed({
+  model: registry.textEmbeddingModel('openai:text-embedding-3-small'),
+  value: 'sunny day at the beach',
+});
+```

content/docs/07-reference/ai-sdk-core/40-provider-registry.mdx

@@ -13,40 +13,6 @@ in a central place and access the models through simple string ids.
 `createProviderRegistry` lets you create a registry with multiple providers that you
 can access by their ids.
 
-### Setup (Example)
-
-You can create a registry with multiple providers and models using `createProviderRegistry`.
-
-```ts
-import { anthropic } from '@ai-sdk/anthropic';
-import { createOpenAI } from '@ai-sdk/openai';
-import { experimental_createProviderRegistry as createProviderRegistry } from 'ai';
-
-export const registry = createProviderRegistry({
-  // register provider with prefix and default setup:
-  anthropic,
-
-  // register provider with prefix and custom setup:
-  openai: createOpenAI({
-    apiKey: process.env.OPENAI_API_KEY,
-  }),
-});
-```
-
-### Usage (Example)
-
-You can access language models by using the `languageModel` method on the registry.
-The provider id will become the prefix of the model id: `providerId:modelId`.
-
-```ts highlight={"4"}
-import { generateText } from 'ai';
-
-const { text } = await generateText({
-  model: registry.languageModel('openai:gpt-4-turbo'),
-  prompt: 'Invent a new holiday and describe its traditions.',
-});
-```
-
 ## Import
 
 <Snippet
@@ -64,7 +30,7 @@ Registers a language model provider with a given id.
   content={[
     {
       name: 'providers',
-      type: 'Record<string, (id: string) => LanguageModel>',
+      type: 'Record<string, { languageModel: (id: string) => LanguageModel; textEmbedding: (id: string) => EmbeddingModel<string> }>',
       description: `The unique identifier for the provider. It should be unique within the registry.`,
     },
   ]}
@@ -81,5 +47,62 @@ The `experimental_createProviderRegistry` function returns a `experimental_Provi
       type: '(id: string) => LanguageModel',
       description: `A function that returns a language model by its id (format: providerId:modelId)`,
     },
+    {
+      name: 'textEmbeddingModel',
+      type: '(id: string) => EmbeddingModel<string>',
+      description: `A function that returns a text embedding model by its id (format: providerId:modelId)`,
+    },
   ]}
 />
+
+## Examples
+
+### Setup
+
+You can create a registry with multiple providers and models using `createProviderRegistry`.
+
+```ts
+import { anthropic } from '@ai-sdk/anthropic';
+import { createOpenAI } from '@ai-sdk/openai';
+import { experimental_createProviderRegistry as createProviderRegistry } from 'ai';
+
+export const registry = createProviderRegistry({
+  // register provider with prefix and default setup:
+  anthropic,
+
+  // register provider with prefix and custom setup:
+  openai: createOpenAI({
+    apiKey: process.env.OPENAI_API_KEY,
+  }),
+});
+```
+
+### Language models
+
+You can access language models by using the `languageModel` method on the registry.
+The provider id will become the prefix of the model id: `providerId:modelId`.
+
+```ts highlight={"5"}
+import { generateText } from 'ai';
+import { registry } from './registry';
+
+const { text } = await generateText({
+  model: registry.languageModel('openai:gpt-4-turbo'),
+  prompt: 'Invent a new holiday and describe its traditions.',
+});
+```
+
+### Text embedding models
+
+You can access text embedding models by using the `textEmbeddingModel` method on the registry.
+The provider id will become the prefix of the model id: `providerId:modelId`.
+
+```ts highlight={"5"}
+import { embed } from 'ai';
+import { registry } from './registry';
+
+const { embedding } = await embed({
+  model: registry.textEmbeddingModel('openai:text-embedding-3-small'),
+  value: 'sunny day at the beach',
+});
+```

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

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

examples/ai-core/src/registry/setup-registry.ts

@@ -1,4 +1,5 @@
 import { anthropic } from '@ai-sdk/anthropic';
+import { mistral } from '@ai-sdk/mistral';
 import { createOpenAI } from '@ai-sdk/openai';
 import { experimental_createProviderRegistry as createProviderRegistry } from 'ai';
 import dotenv from 'dotenv';
@@ -8,6 +9,7 @@ dotenv.config();
 export const registry = createProviderRegistry({
   // register provider with prefix and default setup:
   anthropic,
+  mistral,
 
   // register provider with prefix and custom setup:
   openai: createOpenAI({

packages/anthropic/src/anthropic-provider.ts

@@ -16,6 +16,14 @@ Creates a model for text generation.
 
   /**
 Creates a model for text generation.
+*/
+  languageModel(
+    modelId: AnthropicMessagesModelId,
+    settings?: AnthropicMessagesSettings,
+  ): AnthropicMessagesLanguageModel;
+
+  /**
+Creates a model for text generation.
 */
   chat(
     modelId: AnthropicMessagesModelId,
@@ -108,6 +116,7 @@ export function createAnthropic(
     return createChatModel(modelId, settings);
   };
 
+  provider.languageModel = createChatModel;
   provider.chat = createChatModel;
   provider.messages = createChatModel;
 

packages/azure/src/azure-openai-provider.ts

@@ -10,6 +10,14 @@ export interface AzureOpenAIProvider {
     settings?: OpenAIChatSettings,
   ): OpenAIChatLanguageModel;
 
+  /**
+Creates an Azure OpenAI chat model for text generation.
+   */
+  languageModel(
+    deploymentId: string,
+    settings?: OpenAIChatSettings,
+  ): OpenAIChatLanguageModel;
+
   /**
 Creates an Azure OpenAI chat model for text generation.
    */
@@ -85,6 +93,7 @@ export function createAzure(
     return createChatModel(deploymentId, settings as OpenAIChatSettings);
   };
 
+  provider.languageModel = createChatModel;
   provider.chat = createChatModel;
 
   return provider as AzureOpenAIProvider;

packages/core/core/registry/no-such-model-error.ts

@@ -1,25 +1,30 @@
 export class NoSuchModelError extends Error {
   readonly modelId: string;
+  readonly modelType: string;
 
   constructor({
     modelId,
-    message = `No such model: ${modelId}`,
+    modelType,
+    message = `No such ${modelType}: ${modelId}`,
   }: {
     modelId: string;
+    modelType: string;
     message?: string;
   }) {
     super(message);
 
     this.name = 'AI_NoSuchModelError';
 
     this.modelId = modelId;
+    this.modelType = modelType;
   }
 
   static isNoSuchModelError(error: unknown): error is NoSuchModelError {
     return (
       error instanceof Error &&
       error.name === 'AI_NoSuchModelError' &&
-      typeof (error as NoSuchModelError).modelId === 'string'
+      typeof (error as NoSuchModelError).modelId === 'string' &&
+      typeof (error as NoSuchModelError).modelType === 'string'
     );
   }
 
@@ -30,6 +35,7 @@ export class NoSuchModelError extends Error {
       stack: this.stack,
 
       modelId: this.modelId,
+      modelType: this.modelType,
     };
   }
 }

packages/core/core/registry/no-such-provider-error.ts

@@ -1,25 +1,30 @@
 export class NoSuchProviderError extends Error {
   readonly providerId: string;
+  readonly availableProviders: string[];
 
   constructor({
     providerId,
-    message = `No such provider: ${providerId}`,
+    availableProviders,
+    message = `No such provider: ${providerId} (available providers: ${availableProviders.join()})`,
   }: {
     providerId: string;
+    availableProviders: string[];
     message?: string;
   }) {
     super(message);
 
     this.name = 'AI_NoSuchProviderError';
 
     this.providerId = providerId;
+    this.availableProviders = availableProviders;
   }
 
   static isNoSuchProviderError(error: unknown): error is NoSuchProviderError {
     return (
       error instanceof Error &&
       error.name === 'AI_NoSuchProviderError' &&
-      typeof (error as NoSuchProviderError).providerId === 'string'
+      typeof (error as NoSuchProviderError).providerId === 'string' &&
+      Array.isArray((error as NoSuchProviderError).availableProviders)
     );
   }
 
@@ -30,6 +35,7 @@ export class NoSuchProviderError extends Error {
       stack: this.stack,
 
       providerId: this.providerId,
+      availableProviders: this.availableProviders,
     };
   }
 }