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

Co-authored-by: Tobias Strebitzer <tobias.strebitzer@magloft.com>

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);

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

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

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

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

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

@@ -5,11 +5,13 @@ import dotenv from 'dotenv';
 dotenv.config();
 
 async function main() {
-  const { embedding } = await embed({
+  const { embedding, usage } = await embed({
     model: azure.embedding('my-embedding-deployment'),
     value: 'sunny day at the beach',
   });
+
   console.log(embedding);
+  console.log(usage);
 }
 
 main().catch(console.error);

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

@@ -5,12 +5,13 @@ import dotenv from 'dotenv';
 dotenv.config();
 
 async function main() {
-  const { embedding } = await embed({
+  const { embedding, usage } = await embed({
     model: mistral.embedding('mistral-embed'),
     value: 'sunny day at the beach',
   });
 
   console.log(embedding);
+  console.log(usage);
 }
 
 main().catch(console.error);

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

@@ -5,12 +5,13 @@ import dotenv from 'dotenv';
 dotenv.config();
 
 async function main() {
-  const { embedding } = await embed({
+  const { embedding, usage } = await embed({
     model: openai.embedding('text-embedding-3-small'),
     value: 'sunny day at the beach',
   });
 
   console.log(embedding);
+  console.log(usage);
 }
 
 main().catch(console.error);

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

@@ -37,19 +37,16 @@ describe('result.embedding', () => {
       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) };
+          switch (callCount++) {
+            case 0:
+              assert.deepStrictEqual(values, testValues.slice(0, 2));
+              return { embeddings: dummyEmbeddings.slice(0, 2) };
+            case 1:
+              assert.deepStrictEqual(values, testValues.slice(2));
+              return { embeddings: dummyEmbeddings.slice(2) };
+            default:
+              throw new Error('Unexpected call');
           }
-
-          if (callCount === 1) {
-            assert.deepStrictEqual(values, testValues.slice(2));
-            callCount++;
-            return { embeddings: dummyEmbeddings.slice(2) };
-          }
-
-          throw new Error('Unexpected call');
         },
       }),
       values: testValues,
@@ -73,6 +70,37 @@ describe('result.values', () => {
   });
 });
 
+describe('result.usage', () => {
+  it('should include usage in the result', async () => {
+    let callCount = 0;
+
+    const result = await embedMany({
+      model: new MockEmbeddingModelV1({
+        maxEmbeddingsPerCall: 2,
+        doEmbed: async () => {
+          switch (callCount++) {
+            case 0:
+              return {
+                embeddings: dummyEmbeddings.slice(0, 2),
+                usage: { tokens: 10 },
+              };
+            case 1:
+              return {
+                embeddings: dummyEmbeddings.slice(2),
+                usage: { tokens: 20 },
+              };
+            default:
+              throw new Error('Unexpected call');
+          }
+        },
+      }),
+      values: testValues,
+    });
+
+    assert.deepStrictEqual(result.usage, { tokens: 30 });
+  });
+});
+
 describe('options.headers', () => {
   it('should set headers', async () => {
     const result = await embedMany({

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

@@ -1,4 +1,5 @@
 import { Embedding, EmbeddingModel } from '../types';
+import { EmbeddingTokenUsage } from '../types/token-usage';
 import { retryWithExponentialBackoff } from '../util/retry-with-exponential-backoff';
 import { splitArray } from '../util/split-array';
 
@@ -66,6 +67,7 @@ Only applicable for HTTP-based providers.
     return new EmbedManyResult({
       values,
       embeddings: modelResponse.embeddings,
+      usage: modelResponse.usage ?? { tokens: NaN },
     });
   }
 
@@ -74,14 +76,17 @@ Only applicable for HTTP-based providers.
 
   // serially embed the chunks:
   const embeddings = [];
+  let tokens = 0;
+
   for (const chunk of valueChunks) {
     const modelResponse = await retry(() =>
       model.doEmbed({ values: chunk, abortSignal, headers }),
     );
     embeddings.push(...modelResponse.embeddings);
+    tokens += modelResponse.usage?.tokens ?? NaN;
   }
 
-  return new EmbedManyResult({ values, embeddings });
+  return new EmbedManyResult({ values, embeddings, usage: { tokens } });
 }
 
 /**
@@ -99,8 +104,18 @@ The embeddings. They are in the same order as the values.
   */
   readonly embeddings: Array<Embedding>;
 
-  constructor(options: { values: Array<VALUE>; embeddings: Array<Embedding> }) {
+  /**
+The embedding token usage.
+  */
+  readonly usage: EmbeddingTokenUsage;
+
+  constructor(options: {
+    values: Array<VALUE>;
+    embeddings: Array<Embedding>;
+    usage: EmbeddingTokenUsage;
+  }) {
     this.values = options.values;
     this.embeddings = options.embeddings;
+    this.usage = options.usage;
   }
 }

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

@@ -34,6 +34,19 @@ describe('result.value', () => {
   });
 });
 
+describe('result.usage', () => {
+  it('should include usage in the result', async () => {
+    const result = await embed({
+      model: new MockEmbeddingModelV1({
+        doEmbed: mockEmbed([testValue], [dummyEmbedding], { tokens: 10 }),
+      }),
+      value: testValue,
+    });
+
+    assert.deepStrictEqual(result.usage, { tokens: 10 });
+  });
+});
+
 describe('options.headers', () => {
   it('should set headers', async () => {
     const result = await embed({

packages/core/core/embed/embed.ts

@@ -1,4 +1,5 @@
 import { Embedding, EmbeddingModel } from '../types';
+import { EmbeddingTokenUsage } from '../types/token-usage';
 import { retryWithExponentialBackoff } from '../util/retry-with-exponential-backoff';
 
 /**
@@ -57,6 +58,7 @@ Only applicable for HTTP-based providers.
   return new EmbedResult({
     value,
     embedding: modelResponse.embeddings[0],
+    usage: modelResponse.usage ?? { tokens: NaN },
     rawResponse: modelResponse.rawResponse,
   });
 }
@@ -76,6 +78,11 @@ The embedding of the value.
   */
   readonly embedding: Embedding;
 
+  /**
+The embedding token usage.
+  */
+  readonly usage: EmbeddingTokenUsage;
+
   /**
 Optional raw response data.
    */
@@ -89,12 +96,12 @@ Response headers.
   constructor(options: {
     value: VALUE;
     embedding: Embedding;
-    rawResponse?: {
-      headers?: Record<string, string>;
-    };
+    usage: EmbeddingTokenUsage;
+    rawResponse?: { headers?: Record<string, string> };
   }) {
     this.value = options.value;
     this.embedding = options.embedding;
+    this.usage = options.usage;
     this.rawResponse = options.rawResponse;
   }
 }

packages/core/core/generate-object/generate-object.ts

@@ -1,17 +1,20 @@
 import { NoObjectGeneratedError } from '@ai-sdk/provider';
 import { safeParseJSON } from '@ai-sdk/provider-utils';
 import { z } from 'zod';
-import { TokenUsage, calculateTokenUsage } from '../generate-text/token-usage';
 import { CallSettings } from '../prompt/call-settings';
 import { convertToLanguageModelPrompt } from '../prompt/convert-to-language-model-prompt';
 import { getValidatedPrompt } from '../prompt/get-validated-prompt';
 import { prepareCallSettings } from '../prompt/prepare-call-settings';
 import { Prompt } from '../prompt/prompt';
 import { CallWarning, FinishReason, LanguageModel, LogProbs } from '../types';
+import {
+  CompletionTokenUsage,
+  calculateCompletionTokenUsage,
+} from '../types/token-usage';
 import { convertZodToJSONSchema } from '../util/convert-zod-to-json-schema';
+import { prepareResponseHeaders } from '../util/prepare-response-headers';
 import { retryWithExponentialBackoff } from '../util/retry-with-exponential-backoff';
 import { injectJsonSchemaIntoSystem } from './inject-json-schema-into-system';
-import { prepareResponseHeaders } from '../util/prepare-response-headers';
 
 /**
 Generate a structured, typed object for a given prompt and schema using a language model.
@@ -99,7 +102,7 @@ Default and recommended: 'auto' (best mode for the model).
 
   let result: string;
   let finishReason: FinishReason;
-  let usage: Parameters<typeof calculateTokenUsage>[0];
+  let usage: Parameters<typeof calculateCompletionTokenUsage>[0];
   let warnings: CallWarning[] | undefined;
   let rawResponse: { headers?: Record<string, string> } | undefined;
   let logprobs: LogProbs | undefined;
@@ -228,7 +231,7 @@ Default and recommended: 'auto' (best mode for the model).
   return new GenerateObjectResult({
     object: parseResult.value,
     finishReason,
-    usage: calculateTokenUsage(usage),
+    usage: calculateCompletionTokenUsage(usage),
     warnings,
     rawResponse,
     logprobs,
@@ -252,7 +255,7 @@ The reason why the generation finished.
   /**
 The token usage of the generated text.
    */
-  readonly usage: TokenUsage;
+  readonly usage: CompletionTokenUsage;
 
   /**
 Warnings from the model provider (e.g. unsupported settings)
@@ -278,7 +281,7 @@ Logprobs for the completion.
   constructor(options: {
     object: T;
     finishReason: FinishReason;
-    usage: TokenUsage;
+    usage: CompletionTokenUsage;
     warnings: CallWarning[] | undefined;
     rawResponse?: {
       headers?: Record<string, string>;

packages/core/core/generate-object/stream-object.ts

@@ -8,23 +8,26 @@ import {
   isDeepEqualData,
   parsePartialJson,
 } from '@ai-sdk/ui-utils';
+import { ServerResponse } from 'http';
 import { z } from 'zod';
-import { TokenUsage, calculateTokenUsage } from '../generate-text/token-usage';
 import { CallSettings } from '../prompt/call-settings';
 import { convertToLanguageModelPrompt } from '../prompt/convert-to-language-model-prompt';
 import { getValidatedPrompt } from '../prompt/get-validated-prompt';
 import { prepareCallSettings } from '../prompt/prepare-call-settings';
 import { Prompt } from '../prompt/prompt';
 import { CallWarning, FinishReason, LanguageModel, LogProbs } from '../types';
+import {
+  CompletionTokenUsage,
+  calculateCompletionTokenUsage,
+} from '../types/token-usage';
 import {
   AsyncIterableStream,
   createAsyncIterableStream,
 } from '../util/async-iterable-stream';
 import { convertZodToJSONSchema } from '../util/convert-zod-to-json-schema';
+import { prepareResponseHeaders } from '../util/prepare-response-headers';
 import { retryWithExponentialBackoff } from '../util/retry-with-exponential-backoff';
 import { injectJsonSchemaIntoSystem } from './inject-json-schema-into-system';
-import { prepareResponseHeaders } from '../util/prepare-response-headers';
-import { ServerResponse } from 'http';
 
 /**
 Generate a structured, typed object for a given prompt and schema using a language model.
@@ -110,7 +113,7 @@ Callback that is called when the LLM response and the final object validation ar
       /**
 The token usage of the generated response.
 */
-      usage: TokenUsage;
+      usage: CompletionTokenUsage;
 
       /**
 The generated object (typed according to the schema). Can be undefined if the final object does not match the schema.
@@ -327,7 +330,7 @@ The generated object (typed according to the schema). Resolved when the response
   /**
 The token usage of the generated response. Resolved when the response is finished.
    */
-  readonly usage: Promise<TokenUsage>;
+  readonly usage: Promise<CompletionTokenUsage>;
 
   /**
 Optional raw response data.
@@ -368,13 +371,15 @@ Response headers.
     });
 
     // initialize usage promise
-    let resolveUsage: (value: TokenUsage | PromiseLike<TokenUsage>) => void;
-    this.usage = new Promise<TokenUsage>(resolve => {
+    let resolveUsage: (
+      value: CompletionTokenUsage | PromiseLike<CompletionTokenUsage>,
+    ) => void;
+    this.usage = new Promise<CompletionTokenUsage>(resolve => {
       resolveUsage = resolve;
     });
 
     // store information for onFinish callback:
-    let usage: TokenUsage | undefined;
+    let usage: CompletionTokenUsage | undefined;
     let object: T | undefined;
     let error: unknown | undefined;
 
@@ -425,7 +430,7 @@ Response headers.
               }
 
               // store usage for promises and onFinish callback:
-              usage = calculateTokenUsage(chunk.usage);
+              usage = calculateCompletionTokenUsage(chunk.usage);
 
               controller.enqueue({ ...chunk, usage });
 

packages/core/core/generate-text/generate-text.ts

@@ -16,8 +16,11 @@ import {
   LanguageModel,
   LogProbs,
 } from '../types';
+import {
+  CompletionTokenUsage,
+  calculateCompletionTokenUsage,
+} from '../types/token-usage';
 import { retryWithExponentialBackoff } from '../util/retry-with-exponential-backoff';
-import { TokenUsage, calculateTokenUsage } from './token-usage';
 import { ToToolCallArray, parseToolCall } from './tool-call';
 import { ToToolResultArray } from './tool-result';
 
@@ -176,7 +179,7 @@ By default, it's set to 0, which will disable the feature.
     toolCalls: currentToolCalls,
     toolResults: currentToolResults,
     finishReason: currentModelResponse.finishReason,
-    usage: calculateTokenUsage(currentModelResponse.usage),
+    usage: calculateCompletionTokenUsage(currentModelResponse.usage),
     warnings: currentModelResponse.warnings,
     rawResponse: currentModelResponse.rawResponse,
     logprobs: currentModelResponse.logprobs,
@@ -243,7 +246,7 @@ The reason why the generation finished.
   /**
 The token usage of the generated text.
    */
-  readonly usage: TokenUsage;
+  readonly usage: CompletionTokenUsage;
 
   /**
 Warnings from the model provider (e.g. unsupported settings)
@@ -280,7 +283,7 @@ Logprobs for the completion.
     toolCalls: ToToolCallArray<TOOLS>;
     toolResults: ToToolResultArray<TOOLS>;
     finishReason: FinishReason;
-    usage: TokenUsage;
+    usage: CompletionTokenUsage;
     warnings: CallWarning[] | undefined;
     rawResponse?: {
       headers?: Record<string, string>;

packages/core/core/generate-text/index.ts

@@ -1,3 +1,2 @@
 export * from './generate-text';
 export * from './stream-text';
-export type { TokenUsage } from './token-usage';

packages/core/core/generate-text/run-tools-transformation.ts

@@ -1,8 +1,8 @@
 import { LanguageModelV1StreamPart, NoSuchToolError } from '@ai-sdk/provider';
 import { generateId } from '@ai-sdk/ui-utils';
 import { CoreTool } from '../tool';
+import { calculateCompletionTokenUsage } from '../types/token-usage';
 import { TextStreamPart } from './stream-text';
-import { calculateTokenUsage } from './token-usage';
 import { parseToolCall } from './tool-call';
 
 export function runToolsTransformation<TOOLS extends Record<string, CoreTool>>({
@@ -131,7 +131,7 @@ export function runToolsTransformation<TOOLS extends Record<string, CoreTool>>({
             type: 'finish',
             finishReason: chunk.finishReason,
             logprobs: chunk.logprobs,
-            usage: calculateTokenUsage(chunk.usage),
+            usage: calculateCompletionTokenUsage(chunk.usage),
           });
           break;
         }

packages/core/core/generate-text/stream-text.ts

@@ -18,14 +18,14 @@ import {
   LanguageModel,
   LogProbs,
 } from '../types';
+import { CompletionTokenUsage } from '../types/token-usage';
 import {
   AsyncIterableStream,
   createAsyncIterableStream,
 } from '../util/async-iterable-stream';
 import { prepareResponseHeaders } from '../util/prepare-response-headers';
 import { retryWithExponentialBackoff } from '../util/retry-with-exponential-backoff';
 import { runToolsTransformation } from './run-tools-transformation';
-import { TokenUsage } from './token-usage';
 import { ToToolCall } from './tool-call';
 import { ToToolResult } from './tool-result';
 
@@ -109,7 +109,7 @@ The reason why the generation finished.
       /**
 The token usage of the generated response.
  */
-      usage: TokenUsage;
+      usage: CompletionTokenUsage;
 
       /**
 The full text that has been generated.
@@ -210,7 +210,7 @@ Warnings from the model provider (e.g. unsupported settings).
   /**
 The token usage of the generated response. Resolved when the response is finished.
    */
-  readonly usage: Promise<TokenUsage>;
+  readonly usage: Promise<CompletionTokenUsage>;
 
   /**
 The reason why the generation finished. Resolved when the response is finished.
@@ -260,8 +260,10 @@ Response headers.
     this.onFinish = onFinish;
 
     // initialize usage promise
-    let resolveUsage: (value: TokenUsage | PromiseLike<TokenUsage>) => void;
-    this.usage = new Promise<TokenUsage>(resolve => {
+    let resolveUsage: (
+      value: CompletionTokenUsage | PromiseLike<CompletionTokenUsage>,
+    ) => void;
+    this.usage = new Promise<CompletionTokenUsage>(resolve => {
       resolveUsage = resolve;
     });
 
@@ -297,7 +299,7 @@ Response headers.
 
     // store information for onFinish callback:
     let finishReason: FinishReason | undefined;
-    let usage: TokenUsage | undefined;
+    let usage: CompletionTokenUsage | undefined;
     let text = '';
     const toolCalls: ToToolCall<TOOLS>[] = [];
     const toolResults: ToToolResult<TOOLS>[] = [];

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

@@ -1,5 +1,6 @@
 import { EmbeddingModelV1 } from '@ai-sdk/provider';
 import { Embedding } from '../types';
+import { EmbeddingTokenUsage } from '../types/token-usage';
 
 export class MockEmbeddingModelV1<VALUE> implements EmbeddingModelV1<VALUE> {
   readonly specificationVersion = 'v1';
@@ -35,10 +36,11 @@ export class MockEmbeddingModelV1<VALUE> implements EmbeddingModelV1<VALUE> {
 export function mockEmbed<VALUE>(
   expectedValues: Array<VALUE>,
   embeddings: Array<Embedding>,
+  usage?: EmbeddingTokenUsage,
 ): EmbeddingModelV1<VALUE>['doEmbed'] {
   return async ({ values }) => {
     assert.deepStrictEqual(expectedValues, values);
-    return { embeddings };
+    return { embeddings, usage };
   };
 }
 

packages/core/core/types/index.ts

@@ -1,3 +1,13 @@
+import type { CompletionTokenUsage as CompletionTokenUsageOriginal } from './token-usage';
+
 export * from './embedding-model';
 export * from './errors';
 export * from './language-model';
+
+/**
+ * @deprecated Use CompletionTokenUsage instead.
+ */
+export type TokenUsage = CompletionTokenUsageOriginal;
+export type CompletionTokenUsage = CompletionTokenUsageOriginal;
+
+export type { EmbeddingTokenUsage } from './token-usage';

packages/core/core/generate-text/token-usage.ts renamed to packages/core/core/types/token-usage.ts

@@ -1,7 +1,7 @@
 /**
 Represents the number of tokens used in a prompt and completion.
  */
-export type TokenUsage = {
+export type CompletionTokenUsage = {
   /**
 The number of tokens used in the prompt
    */
@@ -18,10 +18,20 @@ The total number of tokens used (promptTokens + completionTokens).
   totalTokens: number;
 };
 
-export function calculateTokenUsage(usage: {
+/**
+Represents the number of tokens used in an embedding.
+ */
+export type EmbeddingTokenUsage = {
+  /**
+The number of tokens used in the embedding.
+   */
+  tokens: number;
+};
+
+export function calculateCompletionTokenUsage(usage: {
   promptTokens: number;
   completionTokens: number;
-}): TokenUsage {
+}): CompletionTokenUsage {
   return {
     promptTokens: usage.promptTokens,
     completionTokens: usage.completionTokens,

packages/core/rsc/stream-ui/stream-ui.tsx

@@ -14,13 +14,13 @@ import { prepareCallSettings } from '../../core/prompt/prepare-call-settings';
 import { prepareToolsAndToolChoice } from '../../core/prompt/prepare-tools-and-tool-choice';
 import { Prompt } from '../../core/prompt/prompt';
 import { CallWarning, CoreToolChoice, FinishReason } from '../../core/types';
+import {
+  CompletionTokenUsage,
+  calculateCompletionTokenUsage,
+} from '../../core/types/token-usage';
 import { retryWithExponentialBackoff } from '../../core/util/retry-with-exponential-backoff';
 import { createStreamableUI } from '../streamable';
 import { createResolvablePromise } from '../utils';
-import {
-  TokenUsage,
-  calculateTokenUsage,
-} from '../../core/generate-text/token-usage';
 
 type Streamable = ReactNode | Promise<ReactNode>;
 
@@ -123,7 +123,7 @@ export async function streamUI<
       /**
        * The token usage of the generated response.
        */
-      usage: TokenUsage;
+      usage: CompletionTokenUsage;
       /**
        * The final ui node that was generated.
        */
@@ -350,7 +350,7 @@ export async function streamUI<
           case 'finish': {
             onFinish?.({
               finishReason: value.finishReason,
-              usage: calculateTokenUsage(value.usage),
+              usage: calculateCompletionTokenUsage(value.usage),
               value: ui.value,
               warnings: result.warnings,
               rawResponse: result.rawResponse,

packages/mistral/src/mistral-embedding-model.test.ts

@@ -18,8 +18,10 @@ describe('doEmbed', () => {
 
   function prepareJsonResponse({
     embeddings = dummyEmbeddings,
+    usage = { prompt_tokens: 8, total_tokens: 8 },
   }: {
     embeddings?: EmbeddingModelV1Embedding[];
+    usage?: { prompt_tokens: number; total_tokens: number };
   } = {}) {
     server.responseBodyJson = {
       id: 'b322cfc2b9d34e2f8e14fc99874faee5',
@@ -30,7 +32,7 @@ describe('doEmbed', () => {
         index: i,
       })),
       model: 'mistral-embed',
-      usage: { prompt_tokens: 8, total_tokens: 8, completion_tokens: 0 },
+      usage,
     };
   }
 
@@ -42,6 +44,16 @@ describe('doEmbed', () => {
     expect(embeddings).toStrictEqual(dummyEmbeddings);
   });
 
+  it('should extract usage', async () => {
+    prepareJsonResponse({
+      usage: { prompt_tokens: 20, total_tokens: 20 },
+    });
+
+    const { usage } = await model.doEmbed({ values: testValues });
+
+    expect(usage).toStrictEqual({ tokens: 20 });
+  });
+
   it('should expose the raw response headers', async () => {
     prepareJsonResponse();
 
@@ -53,7 +65,7 @@ describe('doEmbed', () => {
 
     expect(rawResponse?.headers).toStrictEqual({
       // default headers:
-      'content-length': '289',
+      'content-length': '267',
       'content-type': 'application/json',
 
       // custom header

packages/mistral/src/mistral-embedding-model.ts

@@ -86,6 +86,9 @@ export class MistralEmbeddingModel implements EmbeddingModelV1<string> {
 
     return {
       embeddings: response.data.map(item => item.embedding),
+      usage: response.usage
+        ? { tokens: response.usage.prompt_tokens }
+        : undefined,
       rawResponse: { headers: responseHeaders },
     };
   }
@@ -94,9 +97,6 @@ export class MistralEmbeddingModel implements EmbeddingModelV1<string> {
 // minimal version of the schema, focussed on what is needed for the implementation
 // this approach limits breakages when the API changes and increases efficiency
 const MistralTextEmbeddingResponseSchema = z.object({
-  data: z.array(
-    z.object({
-      embedding: z.array(z.number()),
-    }),
-  ),
+  data: z.array(z.object({ embedding: z.array(z.number()) })),
+  usage: z.object({ prompt_tokens: z.number() }).nullish(),
 });

packages/openai/src/openai-embedding-model.test.ts

@@ -18,8 +18,10 @@ describe('doEmbed', () => {
 
   function prepareJsonResponse({
     embeddings = dummyEmbeddings,
+    usage = { prompt_tokens: 8, total_tokens: 8 },
   }: {
     embeddings?: EmbeddingModelV1Embedding[];
+    usage?: { prompt_tokens: number; total_tokens: number };
   } = {}) {
     server.responseBodyJson = {
       object: 'list',
@@ -29,7 +31,7 @@ describe('doEmbed', () => {
         embedding,
       })),
       model: 'text-embedding-3-large',
-      usage: { prompt_tokens: 8, total_tokens: 8 },
+      usage,
     };
   }
 
@@ -60,6 +62,16 @@ describe('doEmbed', () => {
     });
   });
 
+  it('should extract usage', async () => {
+    prepareJsonResponse({
+      usage: { prompt_tokens: 20, total_tokens: 20 },
+    });
+
+    const { usage } = await model.doEmbed({ values: testValues });
+
+    expect(usage).toStrictEqual({ tokens: 20 });
+  });
+
   it('should pass the model and the values', async () => {
     prepareJsonResponse();
 

packages/openai/src/openai-embedding-model.ts

@@ -89,6 +89,9 @@ export class OpenAIEmbeddingModel implements EmbeddingModelV1<string> {
 
     return {
       embeddings: response.data.map(item => item.embedding),
+      usage: response.usage
+        ? { tokens: response.usage.prompt_tokens }
+        : undefined,
       rawResponse: { headers: responseHeaders },
     };
   }
@@ -97,9 +100,6 @@ export class OpenAIEmbeddingModel implements EmbeddingModelV1<string> {
 // minimal version of the schema, focussed on what is needed for the implementation
 // this approach limits breakages when the API changes and increases efficiency
 const openaiTextEmbeddingResponseSchema = z.object({
-  data: z.array(
-    z.object({
-      embedding: z.array(z.number()),
-    }),
-  ),
+  data: z.array(z.object({ embedding: z.array(z.number()) })),
+  usage: z.object({ prompt_tokens: z.number() }).nullish(),
 });

packages/provider/src/embedding-model/v1/embedding-model-v1.ts

@@ -66,6 +66,11 @@ Generated embeddings. They are in the same order as the input values.
      */
     embeddings: Array<EmbeddingModelV1Embedding>;
 
+    /**
+Token usage. We only have input tokens for embeddings.
+    */
+    usage?: { tokens: number };
+
     /**
 Optional raw response information for debugging purposes.
      */