feat (ai/core): add telemetry support for embed function. (#2385)

.changeset/quick-ears-tie.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+feat (ai/core): add telemetry support for embed function.

content/docs/03-ai-sdk-core/60-telemetry.mdx

@@ -56,12 +56,12 @@ const result = await generateText({
 `generateText` records 3 types of spans:
 
 - `ai.generateText`: the full length of the generateText call. It contains 1 or more `ai.generateText.doGenerate` spans.
-  It contains the [basic span information](#basic-span-information) and the following attributes:
+  It contains the [basic LLM span information](#basic-llm-span-information) and the following attributes:
   - `operation.name`: `ai.generateText`
   - `ai.prompt`: the prompt that was used when calling `generateText`
   - `ai.settings.maxToolRoundtrips`: the maximum number of tool roundtrips that were set
 - `ai.generateText.doGenerate`: a provider doGenerate call. It can contain `ai.toolCall` spans.
-  It contains the [basic span information](#basic-span-information) and the following attributes:
+  It contains the [basic LLM span information](#basic-llm-span-information) and the following attributes:
   - `operation.name`: `ai.generateText`
   - `ai.prompt.format`: the format of the prompt
   - `ai.prompt.messages`: the messages that were passed into the provider
@@ -72,13 +72,13 @@ const result = await generateText({
 `streamText` records 3 types of spans:
 
 - `ai.streamText`: the full length of the streamText call. It contains a `ai.streamText.doStream` span.
-  It contains the [basic span information](#basic-span-information) and the following attributes:
+  It contains the [basic LLM span information](#basic-llm-span-information) and the following attributes:
   - `operation.name`: `ai.streamText`
   - `ai.prompt`: the prompt that was used when calling `streamText`
 - `ai.streamText.doStream`: a provider doStream call.
   This span contains an `ai.stream.firstChunk` event that is emitted when the first chunk of the stream is received.
   The `doStream` span can also contain `ai.toolCall` spans.
-  It contains the [basic span information](#basic-span-information) and the following attributes:
+  It contains the [basic LLM span information](#basic-llm-span-information) and the following attributes:
   - `operation.name`: `ai.streamText`
   - `ai.prompt.format`: the format of the prompt
   - `ai.prompt.messages`: the messages that were passed into the provider
@@ -89,22 +89,38 @@ const result = await generateText({
 `generateObject` records 2 types of spans:
 
 - `ai.generateObject`: the full length of the generateObject call. It contains 1 or more `ai.generateObject.doGenerate` spans.
-  It contains the [basic span information](#basic-span-information) and the following attributes:
+  It contains the [basic LLM span information](#basic-llm-span-information) and the following attributes:
   - `operation.name`: `ai.generateObject`
   - `ai.prompt`: the prompt that was used when calling `generateObject`
   - `ai.settings.mode`: the object generation mode
 - `ai.generateObject.doGenerate`: a provider doGenerate call.
-  It contains the [basic span information](#basic-span-information) and the following attributes:
-  - `operation.name`: `ai.generateText`
+  It contains the [basic LLM span information](#basic-llm-span-information) and the following attributes:
+  - `operation.name`: `ai.generateObject`
   - `ai.prompt.format`: the format of the prompt
   - `ai.prompt.messages`: the messages that were passed into the provider
   - `ai.settings.mode`: the object generation mode
 
+### embed function
+
+`embed` records 2 types of spans:
+
+- `ai.embed`: the full length of the embed call. It contains 1 `ai.embed.doEmbed` spans.
+  It contains the [basic embedding span information](#basic-embedding-span-information) and the following attributes:
+  - `operation.name`: `ai.embed`
+  - `ai.value`: the value that was passed into the `embed` function
+  - `ai.embedding`: a JSON-stringified embedding
+- `ai.embed.doEmbed`: a provider doEmbed call.
+  It contains the [basic embedding span information](#basic-embedding-span-information) and the following attributes:
+  - `operation.name`: `ai.embed`
+  - `ai.values`: the values that were passed into the provider (array)
+  - `ai.embeddings`: an array of JSON-stringified embeddings
+
 ## Span Details
 
-### Basic span information
+### Basic LLM span information
 
-Many spans (`ai.generateText`, `ai.generateText.doGenerate`, `ai.streamText`, `ai.streamText.doStream`) contain the following attributes:
+Many spans that use LLMs (`ai.generateText`, `ai.generateText.doGenerate`, `ai.streamText`, `ai.streamText.doStream`,
+`ai.generateObject`, `ai.generateObject.doGenerate`) contain the following attributes:
 
 - `ai.finishReason`: the reason why the generation finished
 - `ai.model.id`: the id of the model
@@ -119,6 +135,20 @@ Many spans (`ai.generateText`, `ai.generateText.doGenerate`, `ai.streamText`, `a
 - `ai.usage.promptTokens`: the number of prompt tokens that were used
 - `resource.name`: the functionId that was set through `telemetry.functionId`
 
+### Basic embedding span information
+
+Many spans that use embedding models (`ai.embed`, `ai.embed.doEmbed`) contain the following attributes:
+
+- `ai.finishReason`: the reason why the generation finished
+- `ai.model.id`: the id of the model
+- `ai.model.provider`: the provider of the model
+- `ai.request.headers.*`: the request headers that were passed in through `headers`
+- `ai.settings.maxRetries`: the maximum number of retries that were set
+- `ai.telemetry.functionId`: the functionId that was set through `telemetry.functionId`
+- `ai.telemetry.metadata.*`: the metadata that was passed in through `telemetry.metadata`
+- `ai.usage.tokens`: the number of tokens that were used
+- `resource.name`: the functionId that was set through `telemetry.functionId`
+
 ### Tool call spans
 
 Tool call spans (`ai.toolCall`) contain the following attributes:

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

@@ -60,6 +60,40 @@ const { embedding } = await embed({
       description:
         'Additional HTTP headers to be sent with the request. Only applicable for HTTP-based providers.',
     },
+    {
+      name: 'experimental_telemetry',
+      type: 'TelemetrySettings',
+      isOptional: true,
+      description: 'Telemetry configuration. Experimental feature.',
+      properties: [
+        {
+          type: 'TelemetrySettings',
+          parameters: [
+            {
+              name: 'isEnabled',
+              type: 'boolean',
+              isOptional: true,
+              description:
+                'Enable or disable telemetry. Disabled by default while experimental.',
+            },
+            {
+              name: 'functionId',
+              type: 'string',
+              isOptional: true,
+              description:
+                'Identifier for this function. Used to group telemetry data by function.',
+            },
+            {
+              name: 'metadata',
+              isOptional: true,
+              type: 'Record<string, string | number | boolean | Array<null | undefined | string> | Array<null | undefined | number> | Array<null | undefined | boolean>>',
+              description:
+                'Additional information to include in the telemetry data.',
+            },
+          ],
+        },
+      ],
+    },
   ]}
 />
 

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

@@ -1,8 +1,10 @@
 import assert from 'node:assert';
+import { setTestTracer } from '../telemetry/get-tracer';
 import {
   MockEmbeddingModelV1,
   mockEmbed,
 } from '../test/mock-embedding-model-v1';
+import { MockTracer } from '../test/mock-tracer';
 import { embed } from './embed';
 
 const dummyEmbedding = [0.1, 0.2, 0.3];
@@ -66,3 +68,81 @@ describe('options.headers', () => {
     assert.deepStrictEqual(result.embedding, dummyEmbedding);
   });
 });
+
+describe('telemetry', () => {
+  let tracer: MockTracer;
+
+  beforeEach(() => {
+    tracer = new MockTracer();
+    setTestTracer(tracer);
+  });
+
+  afterEach(() => {
+    setTestTracer(undefined);
+  });
+
+  it('should not record any telemetry data when not explicitly enabled', async () => {
+    await embed({
+      model: new MockEmbeddingModelV1({
+        doEmbed: mockEmbed([testValue], [dummyEmbedding]),
+      }),
+      value: testValue,
+    });
+
+    assert.deepStrictEqual(tracer.jsonSpans, []);
+  });
+
+  it('should record telemetry data when enabled with mode "json"', async () => {
+    await embed({
+      model: new MockEmbeddingModelV1({
+        doEmbed: mockEmbed([testValue], [dummyEmbedding], { tokens: 10 }),
+      }),
+      value: testValue,
+      experimental_telemetry: {
+        isEnabled: true,
+        functionId: 'test-function-id',
+        metadata: {
+          test1: 'value1',
+          test2: false,
+        },
+      },
+    });
+
+    assert.deepStrictEqual(tracer.jsonSpans, [
+      {
+        attributes: {
+          'ai.model.id': 'mock-model-id',
+          'ai.model.provider': 'mock-provider',
+          'ai.settings.maxRetries': undefined,
+          'ai.telemetry.functionId': 'test-function-id',
+          'ai.telemetry.metadata.test1': 'value1',
+          'ai.telemetry.metadata.test2': false,
+          'ai.value': '"sunny day at the beach"',
+          'ai.embedding': '[0.1,0.2,0.3]',
+          'ai.usage.tokens': 10,
+          'operation.name': 'ai.embed',
+          'resource.name': 'test-function-id',
+        },
+        events: [],
+        name: 'ai.embed',
+      },
+      {
+        attributes: {
+          'ai.embeddings': ['[0.1,0.2,0.3]'],
+          'ai.model.id': 'mock-model-id',
+          'ai.model.provider': 'mock-provider',
+          'ai.settings.maxRetries': undefined,
+          'ai.telemetry.functionId': 'test-function-id',
+          'ai.telemetry.metadata.test1': 'value1',
+          'ai.telemetry.metadata.test2': false,
+          'ai.usage.tokens': 10,
+          'ai.values': ['"sunny day at the beach"'],
+          'operation.name': 'ai.embed',
+          'resource.name': 'test-function-id',
+        },
+        events: [],
+        name: 'ai.embed.doEmbed',
+      },
+    ]);
+  });
+});

packages/core/core/embed/embed.ts

@@ -1,3 +1,7 @@
+import { getBaseTelemetryAttributes } from '../telemetry/get-base-telemetry-attributes';
+import { getTracer } from '../telemetry/get-tracer';
+import { recordSpan } from '../telemetry/record-span';
+import { TelemetrySettings } from '../telemetry/telemetry-settings';
 import { EmbeddingModel } from '../types';
 import { retryWithExponentialBackoff } from '../util/retry-with-exponential-backoff';
 import { EmbedResult } from './embed-result';
@@ -20,6 +24,7 @@ export async function embed<VALUE>({
   maxRetries,
   abortSignal,
   headers,
+  experimental_telemetry: telemetry,
 }: {
   /**
 The embedding model to use.
@@ -48,18 +53,76 @@ Additional headers to include in the request.
 Only applicable for HTTP-based providers.
  */
   headers?: Record<string, string>;
+
+  /**
+   * Optional telemetry configuration (experimental).
+   */
+  experimental_telemetry?: TelemetrySettings;
 }): Promise<EmbedResult<VALUE>> {
-  const retry = retryWithExponentialBackoff({ maxRetries });
+  const baseTelemetryAttributes = getBaseTelemetryAttributes({
+    operationName: 'ai.embed',
+    model,
+    telemetry,
+    headers,
+    settings: { maxRetries },
+  });
+
+  const tracer = getTracer({ isEnabled: telemetry?.isEnabled ?? false });
+
+  return recordSpan({
+    name: 'ai.embed',
+    attributes: {
+      ...baseTelemetryAttributes,
+      // specific settings that only make sense on the outer level:
+      'ai.value': JSON.stringify(value),
+    },
+    tracer,
+    fn: async span => {
+      const retry = retryWithExponentialBackoff({ maxRetries });
+
+      const { embedding, usage, rawResponse } = await retry(() =>
+        // nested spans to align with the embedMany telemetry data:
+        recordSpan({
+          name: 'ai.embed.doEmbed',
+          attributes: {
+            ...baseTelemetryAttributes,
+            // specific settings that only make sense on the outer level:
+            'ai.values': [JSON.stringify(value)],
+          },
+          tracer,
+          fn: async doEmbedSpan => {
+            const modelResponse = await model.doEmbed({
+              values: [value],
+              abortSignal,
+              headers,
+            });
+
+            const embedding = modelResponse.embeddings[0];
+            const usage = modelResponse.usage ?? { tokens: NaN };
+
+            doEmbedSpan.setAttributes({
+              'ai.embeddings': modelResponse.embeddings.map(embedding =>
+                JSON.stringify(embedding),
+              ),
+              'ai.usage.tokens': usage.tokens,
+            });
+
+            return {
+              embedding,
+              usage,
+              rawResponse: modelResponse.rawResponse,
+            };
+          },
+        }),
+      );
 
-  const modelResponse = await retry(() =>
-    model.doEmbed({ values: [value], abortSignal, headers }),
-  );
+      span.setAttributes({
+        'ai.embedding': JSON.stringify(embedding),
+        'ai.usage.tokens': usage.tokens,
+      });
 
-  return new DefaultEmbedResult({
-    value,
-    embedding: modelResponse.embeddings[0],
-    usage: modelResponse.usage ?? { tokens: NaN },
-    rawResponse: modelResponse.rawResponse,
+      return new DefaultEmbedResult({ value, embedding, usage, rawResponse });
+    },
   });
 }
 

packages/core/core/telemetry/get-base-telemetry-attributes.ts

@@ -1,6 +1,5 @@
 import { Attributes } from '@opentelemetry/api';
 import { CallSettings } from '../prompt/call-settings';
-import { LanguageModel } from '../types/language-model';
 import { TelemetrySettings } from './telemetry-settings';
 
 export function getBaseTelemetryAttributes({
@@ -11,7 +10,7 @@ export function getBaseTelemetryAttributes({
   headers,
 }: {
   operationName: string;
-  model: LanguageModel;
+  model: { modelId: string; provider: string };
   settings: Omit<CallSettings, 'abortSignal' | 'headers'>;
   telemetry: TelemetrySettings | undefined;
   headers: Record<string, string | undefined> | undefined;