feat(core): Instrument Google GenAI embeddings API (embedContent)

nicohrubec · claude · nicohrubec · commit 3af75a0d64b0 · 2026-03-26T19:53:33.000+01:00
Add embeddings support to the Google GenAI integration using the new
method registry architecture. Records embeddings input when PII
recording is enabled, and skips response attribute extraction for
embeddings responses.

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/dev-packages/browser-integration-tests/suites/tracing/ai-providers/google-genai/mocks.js b/dev-packages/browser-integration-tests/suites/tracing/ai-providers/google-genai/mocks.js
@@ -39,6 +39,24 @@ export class MockGoogleGenAI {
           },
         };
       },
+      embedContent: async (...args) => {
+        const params = args[0];
+        await new Promise(resolve => setTimeout(resolve, 10));
+
+        if (params.model === 'error-model') {
+          const error = new Error('Model not found');
+          error.status = 404;
+          throw error;
+        }
+
+        return {
+          embeddings: [
+            {
+              values: [0.1, 0.2, 0.3, 0.4, 0.5],
+            },
+          ],
+        };
+      },
       generateContentStream: async () => {
         // Return a promise that resolves to an async generator
         return (async function* () {
diff --git a/dev-packages/browser-integration-tests/suites/tracing/ai-providers/google-genai/subject.js b/dev-packages/browser-integration-tests/suites/tracing/ai-providers/google-genai/subject.js
@@ -30,3 +30,11 @@ const response = await chat.sendMessage({
 });
 
 console.log('Received response', response);
+
+// Test embedContent
+const embedResponse = await client.models.embedContent({
+  model: 'text-embedding-004',
+  contents: 'Hello world',
+});
+
+console.log('Received embed response', embedResponse);
diff --git a/dev-packages/node-integration-tests/suites/tracing/google-genai/scenario-embeddings.mjs b/dev-packages/node-integration-tests/suites/tracing/google-genai/scenario-embeddings.mjs
@@ -0,0 +1,77 @@
+import { GoogleGenAI } from '@google/genai';
+import * as Sentry from '@sentry/node';
+import express from 'express';
+
+function startMockGoogleGenAIServer() {
+  const app = express();
+  app.use(express.json());
+
+  app.post('/v1beta/models/:model\\:batchEmbedContents', (req, res) => {
+    const model = req.params.model;
+
+    if (model === 'error-model') {
+      res.status(404).set('x-request-id', 'mock-request-123').end('Model not found');
+      return;
+    }
+
+    res.send({
+      embeddings: [
+        {
+          values: [0.1, 0.2, 0.3, 0.4, 0.5],
+        },
+      ],
+    });
+  });
+
+  return new Promise(resolve => {
+    const server = app.listen(0, () => {
+      resolve(server);
+    });
+  });
+}
+
+async function run() {
+  const server = await startMockGoogleGenAIServer();
+
+  await Sentry.startSpan({ op: 'function', name: 'main' }, async () => {
+    const client = new GoogleGenAI({
+      apiKey: 'mock-api-key',
+      httpOptions: { baseUrl: `http://localhost:${server.address().port}` },
+    });
+
+    // Test 1: Basic embedContent with string contents
+    await client.models.embedContent({
+      model: 'text-embedding-004',
+      contents: 'What is the capital of France?',
+    });
+
+    // Test 2: Error handling
+    try {
+      await client.models.embedContent({
+        model: 'error-model',
+        contents: 'This will fail',
+      });
+    } catch {
+      // Expected error
+    }
+
+    // Test 3: embedContent with array contents
+    await client.models.embedContent({
+      model: 'text-embedding-004',
+      contents: [
+        {
+          role: 'user',
+          parts: [{ text: 'First input text' }],
+        },
+        {
+          role: 'user',
+          parts: [{ text: 'Second input text' }],
+        },
+      ],
+    });
+  });
+
+  server.close();
+}
+
+run();
diff --git a/dev-packages/node-integration-tests/suites/tracing/google-genai/test.ts b/dev-packages/node-integration-tests/suites/tracing/google-genai/test.ts
@@ -1,6 +1,7 @@
 import { SEMANTIC_ATTRIBUTE_SENTRY_OP, SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '@sentry/core';
 import { afterAll, describe, expect } from 'vitest';
 import {
+  GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE,
   GEN_AI_INPUT_MESSAGES_ATTRIBUTE,
   GEN_AI_INPUT_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE,
   GEN_AI_OPERATION_NAME_ATTRIBUTE,
@@ -601,4 +602,124 @@ describe('Google GenAI integration', () => {
       });
     },
   );
+
+  const EXPECTED_TRANSACTION_DEFAULT_PII_FALSE_EMBEDDINGS = {
+    transaction: 'main',
+    spans: expect.arrayContaining([
+      // First span - embedContent with string contents
+      expect.objectContaining({
+        data: {
+          [GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'embeddings',
+          [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'gen_ai.embeddings',
+          [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.ai.google_genai',
+          [GEN_AI_SYSTEM_ATTRIBUTE]: 'google_genai',
+          [GEN_AI_REQUEST_MODEL_ATTRIBUTE]: 'text-embedding-004',
+        },
+        description: 'embeddings text-embedding-004',
+        op: 'gen_ai.embeddings',
+        origin: 'auto.ai.google_genai',
+        status: 'ok',
+      }),
+      // Second span - embedContent error model
+      expect.objectContaining({
+        data: {
+          [GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'embeddings',
+          [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'gen_ai.embeddings',
+          [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.ai.google_genai',
+          [GEN_AI_SYSTEM_ATTRIBUTE]: 'google_genai',
+          [GEN_AI_REQUEST_MODEL_ATTRIBUTE]: 'error-model',
+        },
+        description: 'embeddings error-model',
+        op: 'gen_ai.embeddings',
+        origin: 'auto.ai.google_genai',
+        status: 'internal_error',
+      }),
+      // Third span - embedContent with array contents
+      expect.objectContaining({
+        data: {
+          [GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'embeddings',
+          [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'gen_ai.embeddings',
+          [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.ai.google_genai',
+          [GEN_AI_SYSTEM_ATTRIBUTE]: 'google_genai',
+          [GEN_AI_REQUEST_MODEL_ATTRIBUTE]: 'text-embedding-004',
+        },
+        description: 'embeddings text-embedding-004',
+        op: 'gen_ai.embeddings',
+        origin: 'auto.ai.google_genai',
+        status: 'ok',
+      }),
+    ]),
+  };
+
+  const EXPECTED_TRANSACTION_DEFAULT_PII_TRUE_EMBEDDINGS = {
+    transaction: 'main',
+    spans: expect.arrayContaining([
+      // First span - embedContent with PII
+      expect.objectContaining({
+        data: {
+          [GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'embeddings',
+          [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'gen_ai.embeddings',
+          [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.ai.google_genai',
+          [GEN_AI_SYSTEM_ATTRIBUTE]: 'google_genai',
+          [GEN_AI_REQUEST_MODEL_ATTRIBUTE]: 'text-embedding-004',
+          [GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE]: 'What is the capital of France?',
+        },
+        description: 'embeddings text-embedding-004',
+        op: 'gen_ai.embeddings',
+        origin: 'auto.ai.google_genai',
+        status: 'ok',
+      }),
+      // Second span - embedContent error model with PII
+      expect.objectContaining({
+        data: {
+          [GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'embeddings',
+          [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'gen_ai.embeddings',
+          [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.ai.google_genai',
+          [GEN_AI_SYSTEM_ATTRIBUTE]: 'google_genai',
+          [GEN_AI_REQUEST_MODEL_ATTRIBUTE]: 'error-model',
+          [GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE]: 'This will fail',
+        },
+        description: 'embeddings error-model',
+        op: 'gen_ai.embeddings',
+        origin: 'auto.ai.google_genai',
+        status: 'internal_error',
+      }),
+      // Third span - embedContent with array contents and PII
+      expect.objectContaining({
+        data: {
+          [GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'embeddings',
+          [SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'gen_ai.embeddings',
+          [SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.ai.google_genai',
+          [GEN_AI_SYSTEM_ATTRIBUTE]: 'google_genai',
+          [GEN_AI_REQUEST_MODEL_ATTRIBUTE]: 'text-embedding-004',
+          [GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE]:
+            '[{"role":"user","parts":[{"text":"First input text"}]},{"role":"user","parts":[{"text":"Second input text"}]}]',
+        },
+        description: 'embeddings text-embedding-004',
+        op: 'gen_ai.embeddings',
+        origin: 'auto.ai.google_genai',
+        status: 'ok',
+      }),
+    ]),
+  };
+
+  createEsmAndCjsTests(__dirname, 'scenario-embeddings.mjs', 'instrument.mjs', (createRunner, test) => {
+    test('creates google genai embeddings spans with sendDefaultPii: false', async () => {
+      await createRunner()
+        .ignore('event')
+        .expect({ transaction: EXPECTED_TRANSACTION_DEFAULT_PII_FALSE_EMBEDDINGS })
+        .start()
+        .completed();
+    });
+  });
+
+  createEsmAndCjsTests(__dirname, 'scenario-embeddings.mjs', 'instrument-with-pii.mjs', (createRunner, test) => {
+    test('creates google genai embeddings spans with sendDefaultPii: true', async () => {
+      await createRunner()
+        .ignore('event')
+        .expect({ transaction: EXPECTED_TRANSACTION_DEFAULT_PII_TRUE_EMBEDDINGS })
+        .start()
+        .completed();
+    });
+  });
 });
diff --git a/packages/core/src/tracing/google-genai/constants.ts b/packages/core/src/tracing/google-genai/constants.ts
@@ -9,6 +9,7 @@ export const GOOGLE_GENAI_INTEGRATION_NAME = 'Google_GenAI';
 export const GOOGLE_GENAI_METHOD_REGISTRY = {
   'models.generateContent': { operation: 'generate_content' },
   'models.generateContentStream': { operation: 'generate_content', streaming: true },
+  'models.embedContent': { operation: 'embeddings' },
   'chats.create': { operation: 'chat' },
   // chat.* paths are built by createDeepProxy when it proxies the chat instance with CHAT_PATH as base
   'chat.sendMessage': { operation: 'chat' },
diff --git a/packages/core/src/tracing/google-genai/index.ts b/packages/core/src/tracing/google-genai/index.ts
@@ -5,6 +5,7 @@ import { startSpan, startSpanManual } from '../../tracing/trace';
 import type { Span, SpanAttributeValue } from '../../types-hoist/span';
 import { handleCallbackErrors } from '../../utils/handleCallbackErrors';
 import {
+  GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE,
   GEN_AI_INPUT_MESSAGES_ATTRIBUTE,
   GEN_AI_INPUT_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE,
   GEN_AI_OPERATION_NAME_ATTRIBUTE,
@@ -132,7 +133,18 @@ function extractRequestAttributes(
  * This is only recorded if recordInputs is true.
  * Handles different parameter formats for different Google GenAI methods.
  */
-function addPrivateRequestAttributes(span: Span, params: Record<string, unknown>): void {
+function addPrivateRequestAttributes(span: Span, params: Record<string, unknown>, isEmbeddings: boolean): void {
+  if (isEmbeddings) {
+    const contents = params.contents;
+    if (contents != null) {
+      span.setAttribute(
+        GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE,
+        typeof contents === 'string' ? contents : JSON.stringify(contents),
+      );
+    }
+    return;
+  }
+
   const messages: Message[] = [];
 
   // config.systemInstruction: ContentUnion
@@ -252,6 +264,7 @@ function instrumentMethod<T extends unknown[], R>(
   options: GoogleGenAIOptions,
 ): (...args: T) => R | Promise<R> {
   const isSyncCreate = methodPath === CHATS_CREATE_METHOD;
+  const isEmbeddings = instrumentedMethod.operation === 'embeddings';
 
   return new Proxy(originalMethod, {
     apply(target, _, args: T): R | Promise<R> {
@@ -272,7 +285,7 @@ function instrumentMethod<T extends unknown[], R>(
           async (span: Span) => {
             try {
               if (options.recordInputs && params) {
-                addPrivateRequestAttributes(span, params);
+                addPrivateRequestAttributes(span, params, isEmbeddings);
               }
               const stream = await target.apply(context, args);
               return instrumentStream(stream, span, Boolean(options.recordOutputs)) as R;
@@ -300,7 +313,7 @@ function instrumentMethod<T extends unknown[], R>(
         },
         (span: Span) => {
           if (options.recordInputs && params) {
-            addPrivateRequestAttributes(span, params);
+            addPrivateRequestAttributes(span, params, isEmbeddings);
           }
 
           return handleCallbackErrors(
@@ -312,8 +325,8 @@ function instrumentMethod<T extends unknown[], R>(
             },
             () => {},
             result => {
-              // Only add response attributes for content-producing methods, not for chats.create
-              if (!isSyncCreate) {
+              // Only add response attributes for content-producing methods, not for chats.create or embeddings
+              if (!isSyncCreate && !isEmbeddings) {
                 addResponseAttributes(span, result, options.recordOutputs);
               }
             },
