feat: Add embedding instrumentation to AI SDK (#1754)

Resolves
#1625

Author: lforst

e2e/scenarios/ai-sdk-instrumentation/assertions.ts

@@ -189,6 +189,20 @@ function findStreamTrace(events: CapturedLogEvent[]) {
   return { child, operation, parent };
 }
 
+function findEmbedTrace(events: CapturedLogEvent[]) {
+  const operation = findLatestSpan(events, "ai-sdk-embed-operation");
+  const parent = findParentSpan(events, "embed", operation?.span.id);
+
+  return { operation, parent };
+}
+
+function findEmbedManyTrace(events: CapturedLogEvent[]) {
+  const operation = findLatestSpan(events, "ai-sdk-embed-many-operation");
+  const parent = findParentSpan(events, "embedMany", operation?.span.id);
+
+  return { operation, parent };
+}
+
 function findToolTrace(events: CapturedLogEvent[]) {
   const operation = findLatestSpan(events, "ai-sdk-tool-operation");
   const parent = findParentSpan(events, "generateText", operation?.span.id);
@@ -641,6 +655,24 @@ function expectAISDKParentSpan(span: CapturedLogEvent | undefined) {
   ).toBe("string");
 }
 
+function expectEmbeddingTokenMetrics(span: CapturedLogEvent | undefined) {
+  const metrics = span?.metrics as Record<string, unknown> | undefined;
+  const totalTokens = metrics?.tokens;
+  const promptTokens = metrics?.prompt_tokens;
+
+  const tokenMetric =
+    typeof totalTokens === "number"
+      ? totalTokens
+      : typeof promptTokens === "number"
+        ? promptTokens
+        : undefined;
+
+  expect(tokenMetric).toEqual(expect.any(Number));
+  if (typeof tokenMetric === "number") {
+    expect(tokenMetric).toBeGreaterThan(0);
+  }
+}
+
 export function defineAISDKInstrumentationAssertions(options: {
   agentSpanName?: AgentSpanName;
   name: string;
@@ -754,6 +786,50 @@ export function defineAISDKInstrumentationAssertions(options: {
       }
     });
 
+    test("captures trace for embed()", testConfig, () => {
+      const root = findLatestSpan(events, ROOT_NAME);
+      const trace = findEmbedTrace(events);
+
+      expectOperationParentedByRoot(trace.operation, root);
+      expectAISDKParentSpan(trace.parent);
+      expect(operationName(trace.operation)).toBe("embed");
+      expectEmbeddingTokenMetrics(trace.parent);
+      const input = isRecord(trace.parent?.input) ? trace.parent.input : null;
+      expect(typeof input?.value).toBe("string");
+      const output = extractOutputRecord(trace.parent);
+      expect(output).toBeDefined();
+      if (output) {
+        expect(output.embedding).toBeUndefined();
+        expect(output.embedding_length).toEqual(expect.any(Number));
+        expect(output.embedding_length).toBeGreaterThan(0);
+      }
+    });
+
+    test("captures trace for embedMany()", testConfig, () => {
+      const root = findLatestSpan(events, ROOT_NAME);
+      const trace = findEmbedManyTrace(events);
+
+      expectOperationParentedByRoot(trace.operation, root);
+      expectAISDKParentSpan(trace.parent);
+      expect(operationName(trace.operation)).toBe("embed-many");
+      expectEmbeddingTokenMetrics(trace.parent);
+      const input = isRecord(trace.parent?.input) ? trace.parent.input : null;
+      expect(Array.isArray(input?.values)).toBe(true);
+      if (Array.isArray(input?.values)) {
+        expect(input.values.length).toBeGreaterThanOrEqual(2);
+      }
+      const output = extractOutputRecord(trace.parent);
+      expect(output).toBeDefined();
+      if (output) {
+        expect(output.embeddings).toBeUndefined();
+        expect(output.responses).toBeUndefined();
+        expect(output.embedding_count).toEqual(expect.any(Number));
+        expect(output.embedding_count).toBeGreaterThanOrEqual(2);
+        expect(output.embedding_length).toEqual(expect.any(Number));
+        expect(output.embedding_length).toBeGreaterThan(0);
+      }
+    });
+
     if (options.supportsOutputObjectScenario) {
       test(
         "captures Output.object schema on generateText()",

e2e/scenarios/ai-sdk-instrumentation/scenario.impl.mjs

@@ -125,6 +125,9 @@ async function runAISDKInstrumentationScenario(
 ) {
   const instrumentedAI = decorateAI ? decorateAI(options.ai) : options.ai;
   const openaiModel = options.openai("gpt-4o-mini");
+  const openaiEmbeddingModel = options.openai.textEmbeddingModel(
+    "text-embedding-3-small",
+  );
   const sdkMajorVersion = parseMajorVersion(options.sdkVersion);
   const supportsRichInputScenarios = sdkMajorVersion >= 5;
   const outputObject = createOutputObjectIfSupported(options.ai);
@@ -169,6 +172,28 @@ async function runAISDKInstrumentationScenario(
         }
       });
 
+      await runOperation("ai-sdk-embed-operation", "embed", async () => {
+        await instrumentedAI.embed({
+          model: openaiEmbeddingModel,
+          value: "Paris is the capital of France.",
+        });
+      });
+
+      await runOperation(
+        "ai-sdk-embed-many-operation",
+        "embed-many",
+        async () => {
+          await instrumentedAI.embedMany({
+            model: openaiEmbeddingModel,
+            values: [
+              "Paris is in France.",
+              "Berlin is in Germany.",
+              "Vienna is in Austria.",
+            ],
+          });
+        },
+      );
+
       await runOperation("ai-sdk-tool-operation", "tool", async () => {
         const toolRequest = {
           model: openaiModel,

js/src/auto-instrumentations/configs/ai-sdk.test.ts

@@ -0,0 +1,46 @@
+import { describe, expect, it } from "vitest";
+import { aiSDKChannels } from "../../instrumentation/plugins/ai-sdk-channels";
+import { aiSDKConfigs } from "./ai-sdk";
+
+function findConfigsByFunctionName(functionName: string) {
+  return aiSDKConfigs.filter((config) => {
+    if (!("functionQuery" in config)) {
+      return false;
+    }
+    const query = config.functionQuery as { functionName?: unknown };
+    return query.functionName === functionName;
+  });
+}
+
+describe("aiSDKConfigs", () => {
+  it("defines embed channels", () => {
+    expect(aiSDKChannels.embed.channelName).toBe("embed");
+    expect(aiSDKChannels.embedMany.channelName).toBe("embedMany");
+  });
+
+  it("instruments embed() in both ESM and CJS entrypoints", () => {
+    const embedConfigs = findConfigsByFunctionName("embed");
+
+    expect(embedConfigs).toHaveLength(2);
+    expect(embedConfigs.map((config) => config.channelName)).toEqual([
+      aiSDKChannels.embed.channelName,
+      aiSDKChannels.embed.channelName,
+    ]);
+    expect(embedConfigs.map((config) => config.module.filePath).sort()).toEqual(
+      ["dist/index.js", "dist/index.mjs"],
+    );
+  });
+
+  it("instruments embedMany() in both ESM and CJS entrypoints", () => {
+    const embedManyConfigs = findConfigsByFunctionName("embedMany");
+
+    expect(embedManyConfigs).toHaveLength(2);
+    expect(embedManyConfigs.map((config) => config.channelName)).toEqual([
+      aiSDKChannels.embedMany.channelName,
+      aiSDKChannels.embedMany.channelName,
+    ]);
+    expect(
+      embedManyConfigs.map((config) => config.module.filePath).sort(),
+    ).toEqual(["dist/index.js", "dist/index.mjs"]);
+  });
+});

js/src/auto-instrumentations/configs/ai-sdk.ts

@@ -105,6 +105,58 @@ export const aiSDKConfigs: InstrumentationConfig[] = [
     },
   },
 
+  // embed - async function
+  {
+    channelName: aiSDKChannels.embed.channelName,
+    module: {
+      name: "ai",
+      versionRange: ">=3.0.0",
+      filePath: "dist/index.mjs",
+    },
+    functionQuery: {
+      functionName: "embed",
+      kind: "Async",
+    },
+  },
+  {
+    channelName: aiSDKChannels.embed.channelName,
+    module: {
+      name: "ai",
+      versionRange: ">=3.0.0",
+      filePath: "dist/index.js",
+    },
+    functionQuery: {
+      functionName: "embed",
+      kind: "Async",
+    },
+  },
+
+  // embedMany - async function
+  {
+    channelName: aiSDKChannels.embedMany.channelName,
+    module: {
+      name: "ai",
+      versionRange: ">=3.0.0",
+      filePath: "dist/index.mjs",
+    },
+    functionQuery: {
+      functionName: "embedMany",
+      kind: "Async",
+    },
+  },
+  {
+    channelName: aiSDKChannels.embedMany.channelName,
+    module: {
+      name: "ai",
+      versionRange: ">=3.0.0",
+      filePath: "dist/index.js",
+    },
+    functionQuery: {
+      functionName: "embedMany",
+      kind: "Async",
+    },
+  },
+
   // streamObject - async function (v3 only, before the sync refactor in v4)
   {
     channelName: aiSDKChannels.streamObject.channelName,

js/src/instrumentation/plugins/ai-sdk-channels.ts

@@ -3,6 +3,8 @@ import type { ChannelSpanInfo } from "../core/types";
 import type {
   AISDK,
   AISDKCallParams,
+  AISDKEmbedParams,
+  AISDKEmbeddingResult,
   AISDKResult,
 } from "../../vendor-sdk-types/ai-sdk";
 
@@ -69,6 +71,20 @@ export const aiSDKChannels = defineChannels("ai", {
     channelName: "streamObject.sync",
     kind: "sync-stream",
   }),
+  embed: channel<[AISDKEmbedParams], AISDKEmbeddingResult, AISDKChannelContext>(
+    {
+      channelName: "embed",
+      kind: "async",
+    },
+  ),
+  embedMany: channel<
+    [AISDKEmbedParams],
+    AISDKEmbeddingResult,
+    AISDKChannelContext
+  >({
+    channelName: "embedMany",
+    kind: "async",
+  }),
   agentGenerate: channel<
     [AISDKCallParams],
     AISDKStreamResult,

js/src/instrumentation/plugins/ai-sdk-plugin.test.ts

@@ -883,6 +883,53 @@ describe("AI SDK utility functions", () => {
       expect(result).toBeUndefined();
     });
   });
+
+  describe("processAISDKEmbeddingOutput", () => {
+    it("should summarize single embedding length", () => {
+      const output = {
+        embedding: [0.1, 0.2, 0.3, 0.4],
+        usage: {
+          totalTokens: 10,
+        },
+      };
+
+      const result = processAISDKEmbeddingOutput(output, []);
+      expect(result.embedding).toBeUndefined();
+      expect(result.embedding_length).toBe(4);
+      expect(result.usage).toMatchObject({
+        totalTokens: 10,
+      });
+    });
+
+    it("should summarize embedding batches", () => {
+      const output = {
+        embeddings: [
+          [0.1, 0.2, 0.3],
+          [0.4, 0.5, 0.6],
+        ],
+      };
+
+      const result = processAISDKEmbeddingOutput(output, []);
+      expect(result.embeddings).toBeUndefined();
+      expect(result.embedding_count).toBe(2);
+      expect(result.embedding_length).toBe(3);
+    });
+
+    it("should omit non-whitelisted fields like responses", () => {
+      const output = {
+        embeddings: [[0.1, 0.2, 0.3]],
+        response: { body: "too much" },
+        responses: [{ body: "way too much" }],
+        usage: { totalTokens: 8 },
+      };
+
+      const result = processAISDKEmbeddingOutput(output, []);
+      expect(result.response).toBeUndefined();
+      expect(result.responses).toBeUndefined();
+      expect(result.usage).toMatchObject({ totalTokens: 8 });
+      expect(result.embedding_count).toBe(1);
+    });
+  });
 });
 
 // Helper functions exported for testing
@@ -1038,12 +1085,17 @@ function extractGetterValues(obj: any): any {
   const getterNames = [
     "text",
     "object",
+    "value",
+    "values",
+    "embedding",
+    "embeddings",
     "finishReason",
     "usage",
     "totalUsage",
     "toolCalls",
     "toolResults",
     "warnings",
+    "responses",
     "experimental_providerMetadata",
     "providerMetadata",
     "rawResponse",
@@ -1210,3 +1262,45 @@ function processAISDKOutput(output: any, denyOutputPaths: string[]): any {
 
   return omit(merged, denyOutputPaths);
 }
+
+function processAISDKEmbeddingOutput(
+  output: any,
+  denyOutputPaths: string[],
+): any {
+  if (!output || typeof output !== "object") {
+    return output;
+  }
+
+  const processed: Record<string, unknown> = {};
+  const whitelistedFields = [
+    "usage",
+    "totalUsage",
+    "warnings",
+    "providerMetadata",
+    "experimental_providerMetadata",
+  ];
+
+  for (const field of whitelistedFields) {
+    const value = output?.[field];
+    if (value !== undefined && typeof value !== "function") {
+      processed[field] = value;
+    }
+  }
+
+  if (Array.isArray(output?.embedding)) {
+    processed.embedding_length = output.embedding.length;
+  }
+
+  if (Array.isArray(output?.embeddings)) {
+    processed.embedding_count = output.embeddings.length;
+
+    const firstEmbedding = output.embeddings.find((item: unknown) =>
+      Array.isArray(item),
+    );
+    if (Array.isArray(firstEmbedding)) {
+      processed.embedding_length = firstEmbedding.length;
+    }
+  }
+
+  return processed;
+}