feat (ai/core): add usage information for embed and embedMany (#2220)

Co-authored-by: Tobias Strebitzer <[email protected]>

Author: lgrammel

.changeset/cool-snakes-agree.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+chore (ai/core): rename TokenUsage type to CompletionTokenUsage

.changeset/unlucky-owls-admire.md

@@ -0,0 +1,8 @@
+---
+'@ai-sdk/provider': patch
+'@ai-sdk/mistral': patch
+'@ai-sdk/openai': patch
+'ai': patch
+---
+
+feat (ai/core): add token usage to embed and embedMany

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

@@ -69,3 +69,20 @@ console.log(
   `cosine similarity: ${cosineSimilarity(embeddings[0], embeddings[1])}`,
 );
 ```
+
+## Token Usage
+
+Many providers charge based on the number of tokens used to generate embeddings.
+Both `embed` and `embedMany` provide token usage information in the `usage` property of the result object:
+
+```ts highlight={"4,9"}
+import { openai } from '@ai-sdk/openai';
+import { embed } from 'ai';
+
+const { embedding, usage } = await embed({
+  model: openai.embedding('text-embedding-3-small'),
+  value: 'sunny day at the beach',
+});
+
+console.log(usage); // { tokens: 10 }
+```

content/docs/07-reference/ai-sdk-core/01-generate-text.mdx

@@ -367,11 +367,11 @@ console.log(text);
     },
     {
       name: 'usage',
-      type: 'TokenUsage',
+      type: 'CompletionTokenUsage',
       description: 'The token usage of the generated text.',
       properties: [
         {
-          type: 'TokenUsage',
+          type: 'CompletionTokenUsage',
           parameters: [
             {
               name: 'promptTokens',

content/docs/07-reference/ai-sdk-core/02-stream-text.mdx

@@ -433,12 +433,12 @@ for await (const textPart of textStream) {
     },
     {
       name: 'usage',
-      type: 'Promise<TokenUsage>',
+      type: 'Promise<CompletionTokenUsage>',
       description:
         'The token usage of the generated text. Resolved when the response is finished.',
       properties: [
         {
-          type: 'TokenUsage',
+          type: 'CompletionTokenUsage',
           parameters: [
             {
               name: 'promptTokens',

content/docs/07-reference/ai-sdk-core/03-generate-object.mdx

@@ -329,11 +329,11 @@ console.log(JSON.stringify(object, null, 2));
     },
     {
       name: 'usage',
-      type: 'TokenUsage',
+      type: 'CompletionTokenUsage',
       description: 'The token usage of the generated text.',
       properties: [
         {
-          type: 'TokenUsage',
+          type: 'CompletionTokenUsage',
           parameters: [
             {
               name: 'promptTokens',

content/docs/07-reference/ai-sdk-core/04-stream-object.mdx

@@ -325,11 +325,11 @@ for await (const partialObject of partialObjectStream) {
           parameters: [
             {
               name: 'usage',
-              type: 'TokenUsage',
+              type: 'CompletionTokenUsage',
               description: 'The token usage of the generated text.',
               properties: [
                 {
-                  type: 'TokenUsage',
+                  type: 'CompletionTokenUsage',
                   parameters: [
                     {
                       name: 'promptTokens',
@@ -400,12 +400,12 @@ for await (const partialObject of partialObjectStream) {
   content={[
     {
       name: 'usage',
-      type: 'Promise<TokenUsage>',
+      type: 'Promise<CompletionTokenUsage>',
       description:
         'The token usage of the generated text. Resolved when the response is finished.',
       properties: [
         {
-          type: 'TokenUsage',
+          type: 'CompletionTokenUsage',
           parameters: [
             {
               name: 'promptTokens',

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

@@ -77,6 +77,23 @@ const { embedding } = await embed({
       type: 'number[]',
       description: 'The embedding of the value.',
     },
+    {
+      name: 'usage',
+      type: 'EmbeddingTokenUsage',
+      description: 'The token usage for generating the embeddings.',
+      properties: [
+        {
+          type: 'EmbeddingTokenUsage',
+          parameters: [
+            {
+              name: 'tokens',
+              type: 'number',
+              description: 'The total number of input tokens.',
+            },
+          ],
+        },
+      ],
+    },
     {
       name: 'rawResponse',
       type: 'RawResponse',

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

@@ -83,5 +83,22 @@ const { embeddings } = await embedMany({
       type: 'number[][]',
       description: 'The embeddings. They are in the same order as the values.',
     },
+    {
+      name: 'usage',
+      type: 'EmbeddingTokenUsage',
+      description: 'The token usage for generating the embeddings.',
+      properties: [
+        {
+          type: 'EmbeddingTokenUsage',
+          parameters: [
+            {
+              name: 'tokens',
+              type: 'number',
+              description: 'The total number of input tokens.',
+            },
+          ],
+        },
+      ],
+    },
   ]}
 />

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

@@ -5,7 +5,7 @@ import dotenv from 'dotenv';
 dotenv.config();
 
 async function main() {
-  const { embeddings } = await embedMany({
+  const { embeddings, usage } = await embedMany({
     model: azure.embedding('my-embedding-deployment'),
     values: [
       'sunny day at the beach',
@@ -15,6 +15,7 @@ async function main() {
   });
 
   console.log(embeddings);
+  console.log(usage);
 }
 
 main().catch(console.error);