Implement embeddings api support

nicohrubec · nicohrubec · commit dbebf0cd5b87 · 2025-11-17T13:34:43.000+01:00
diff --git a/dev-packages/node-integration-tests/suites/tracing/openai/scenario.mjs b/dev-packages/node-integration-tests/suites/tracing/openai/scenario.mjs
@@ -79,6 +79,13 @@ class MockOpenAI {
       create: async params => {
         await new Promise(resolve => setTimeout(resolve, 10));
 
+        if (params.model === 'error-model') {
+          const error = new Error('Model not found');
+          error.status = 404;
+          error.headers = { 'x-request-id': 'mock-request-123' };
+          throw error;
+        }
+
         return {
           object: 'list',
           data: [
diff --git a/dev-packages/node-integration-tests/suites/tracing/openai/test.ts b/dev-packages/node-integration-tests/suites/tracing/openai/test.ts
@@ -152,7 +152,18 @@ describe('OpenAI integration', () => {
           'sentry.origin': 'auto.ai.openai',
           'gen_ai.system': 'openai',
           'gen_ai.request.model': 'text-embedding-3-small',
+          'gen_ai.request.encoding_format': 'float',
+          'gen_ai.request.dimensions': 1536,
+          'gen_ai.response.model': 'text-embedding-3-small',
+          'gen_ai.usage.input_tokens': 10,
+          'gen_ai.usage.total_tokens': 10,
+          'openai.response.model': 'text-embedding-3-small',
+          'openai.usage.prompt_tokens': 10,
         },
+        description: 'embeddings text-embedding-3-small',
+        op: 'gen_ai.embeddings',
+        origin: 'auto.ai.openai',
+        status: 'ok',
       }),
       // Eighth span - embeddings API error model
       expect.objectContaining({
@@ -163,6 +174,10 @@ describe('OpenAI integration', () => {
           'gen_ai.system': 'openai',
           'gen_ai.request.model': 'error-model',
         },
+        description: 'embeddings error-model',
+        op: 'gen_ai.embeddings',
+        origin: 'auto.ai.openai',
+        status: 'internal_error',
       }),
     ]),
   };
@@ -317,25 +332,42 @@ describe('OpenAI integration', () => {
         origin: 'auto.ai.openai',
         status: 'internal_error',
       }),
-      // Seventh span - embeddings API
+      // Seventh span - embeddings API with PII
       expect.objectContaining({
         data: {
           'gen_ai.operation.name': 'embeddings',
           'sentry.op': 'gen_ai.embeddings',
           'sentry.origin': 'auto.ai.openai',
           'gen_ai.system': 'openai',
           'gen_ai.request.model': 'text-embedding-3-small',
+          'gen_ai.request.encoding_format': 'float',
+          'gen_ai.request.dimensions': 1536,
+          'gen_ai.request.messages': 'Embedding test!',
+          'gen_ai.response.model': 'text-embedding-3-small',
+          'gen_ai.usage.input_tokens': 10,
+          'gen_ai.usage.total_tokens': 10,
+          'openai.response.model': 'text-embedding-3-small',
+          'openai.usage.prompt_tokens': 10,
         },
+        description: 'embeddings text-embedding-3-small',
+        op: 'gen_ai.embeddings',
+        origin: 'auto.ai.openai',
+        status: 'ok',
       }),
-      // Eighth span - embeddings API error model
+      // Eighth span - embeddings API error model with PII
       expect.objectContaining({
         data: {
           'gen_ai.operation.name': 'embeddings',
           'sentry.op': 'gen_ai.embeddings',
           'sentry.origin': 'auto.ai.openai',
           'gen_ai.system': 'openai',
           'gen_ai.request.model': 'error-model',
+          'gen_ai.request.messages': 'Error embedding test!',
         },
+        description: 'embeddings error-model',
+        op: 'gen_ai.embeddings',
+        origin: 'auto.ai.openai',
+        status: 'internal_error',
       }),
     ]),
   };
diff --git a/packages/core/src/tracing/ai/gen-ai-attributes.ts b/packages/core/src/tracing/ai/gen-ai-attributes.ts
@@ -65,6 +65,16 @@ export const GEN_AI_REQUEST_TOP_K_ATTRIBUTE = 'gen_ai.request.top_k';
  */
 export const GEN_AI_REQUEST_STOP_SEQUENCES_ATTRIBUTE = 'gen_ai.request.stop_sequences';
 
+/**
+ * The encoding format for the model request
+ */
+export const GEN_AI_REQUEST_ENCODING_FORMAT_ATTRIBUTE = 'gen_ai.request.encoding_format';
+
+/**
+ * The dimensions for the model request
+ */
+export const GEN_AI_REQUEST_DIMENSIONS_ATTRIBUTE = 'gen_ai.request.dimensions';
+
 /**
  * Array of reasons why the model stopped generating tokens
  */
diff --git a/packages/core/src/tracing/openai/index.ts b/packages/core/src/tracing/openai/index.ts
@@ -7,6 +7,8 @@ import type { Span, SpanAttributeValue } from '../../types-hoist/span';
 import {
   GEN_AI_OPERATION_NAME_ATTRIBUTE,
   GEN_AI_REQUEST_AVAILABLE_TOOLS_ATTRIBUTE,
+  GEN_AI_REQUEST_DIMENSIONS_ATTRIBUTE,
+  GEN_AI_REQUEST_ENCODING_FORMAT_ATTRIBUTE,
   GEN_AI_REQUEST_FREQUENCY_PENALTY_ATTRIBUTE,
   GEN_AI_REQUEST_MESSAGES_ATTRIBUTE,
   GEN_AI_REQUEST_MODEL_ATTRIBUTE,
@@ -15,9 +17,11 @@ import {
   GEN_AI_REQUEST_TEMPERATURE_ATTRIBUTE,
   GEN_AI_REQUEST_TOP_P_ATTRIBUTE,
   GEN_AI_RESPONSE_FINISH_REASONS_ATTRIBUTE,
+  GEN_AI_RESPONSE_MODEL_ATTRIBUTE,
   GEN_AI_RESPONSE_TEXT_ATTRIBUTE,
   GEN_AI_RESPONSE_TOOL_CALLS_ATTRIBUTE,
   GEN_AI_SYSTEM_ATTRIBUTE,
+  OPENAI_RESPONSE_MODEL_ATTRIBUTE,
 } from '../ai/gen-ai-attributes';
 import { getTruncatedJsonString } from '../ai/utils';
 import { OPENAI_INTEGRATION_NAME } from './constants';
@@ -26,6 +30,7 @@ import type {
   ChatCompletionChunk,
   InstrumentedMethod,
   OpenAiChatCompletionObject,
+  OpenAICreateEmbeddingsObject,
   OpenAiIntegration,
   OpenAiOptions,
   OpenAiResponse,
@@ -38,6 +43,7 @@ import {
   getOperationName,
   getSpanOperation,
   isChatCompletionResponse,
+  isEmbeddingsResponse,
   isResponsesApiResponse,
   setCommonResponseAttributes,
   setTokenUsageAttributes,
@@ -82,6 +88,8 @@ function extractRequestAttributes(args: unknown[], methodPath: string): Record<s
       attributes[GEN_AI_REQUEST_FREQUENCY_PENALTY_ATTRIBUTE] = params.frequency_penalty;
     if ('presence_penalty' in params) attributes[GEN_AI_REQUEST_PRESENCE_PENALTY_ATTRIBUTE] = params.presence_penalty;
     if ('stream' in params) attributes[GEN_AI_REQUEST_STREAM_ATTRIBUTE] = params.stream;
+    if ('encoding_format' in params) attributes[GEN_AI_REQUEST_ENCODING_FORMAT_ATTRIBUTE] = params.encoding_format;
+    if ('dimensions' in params) attributes[GEN_AI_REQUEST_DIMENSIONS_ATTRIBUTE] = params.dimensions;
   } else {
     attributes[GEN_AI_REQUEST_MODEL_ATTRIBUTE] = 'unknown';
   }
@@ -166,6 +174,20 @@ function addResponsesApiAttributes(span: Span, response: OpenAIResponseObject, r
   }
 }
 
+/**
+ * Add attributes for Embeddings API responses
+ */
+function addEmbeddingsAttributes(span: Span, response: OpenAICreateEmbeddingsObject): void {
+  span.setAttributes({
+    [OPENAI_RESPONSE_MODEL_ATTRIBUTE]: response.model,
+    [GEN_AI_RESPONSE_MODEL_ATTRIBUTE]: response.model,
+  });
+
+  if (response.usage) {
+    setTokenUsageAttributes(span, response.usage.prompt_tokens, undefined, response.usage.total_tokens);
+  }
+}
+
 /**
  * Add response attributes to spans
  * This currently supports both Chat Completion and Responses API responses
@@ -186,6 +208,8 @@ function addResponseAttributes(span: Span, result: unknown, recordOutputs?: bool
     if (recordOutputs && response.output_text) {
       span.setAttributes({ [GEN_AI_RESPONSE_TEXT_ATTRIBUTE]: response.output_text });
     }
+  } else if (isEmbeddingsResponse(response)) {
+    addEmbeddingsAttributes(span, response);
   }
 }
 
diff --git a/packages/core/src/tracing/openai/types.ts b/packages/core/src/tracing/openai/types.ts
@@ -131,7 +131,29 @@ export interface OpenAIResponseObject {
   metadata: Record<string, unknown>;
 }
 
-export type OpenAiResponse = OpenAiChatCompletionObject | OpenAIResponseObject;
+/**
+ * @see https://platform.openai.com/docs/api-reference/embeddings/object
+ */
+export interface OpenAIEmbeddingsObject {
+  object: 'embedding';
+  embedding: number[];
+  index: number;
+}
+
+/**
+ * @see https://platform.openai.com/docs/api-reference/embeddings/create
+ */
+export interface OpenAICreateEmbeddingsObject {
+  object: 'list';
+  data: OpenAIEmbeddingsObject[];
+  model: string;
+  usage: {
+    prompt_tokens: number;
+    total_tokens: number;
+  };
+}
+
+export type OpenAiResponse = OpenAiChatCompletionObject | OpenAIResponseObject | OpenAICreateEmbeddingsObject;
 
 /**
  * Streaming event types for the Responses API
diff --git a/packages/core/src/tracing/openai/utils.ts b/packages/core/src/tracing/openai/utils.ts
@@ -17,6 +17,7 @@ import type {
   ChatCompletionChunk,
   InstrumentedMethod,
   OpenAiChatCompletionObject,
+  OpenAICreateEmbeddingsObject,
   OpenAIResponseObject,
   ResponseStreamingEvent,
 } from './types';
@@ -83,6 +84,18 @@ export function isResponsesApiResponse(response: unknown): response is OpenAIRes
   );
 }
 
+/**
+ * Check if response is an Embeddings API object
+ */
+export function isEmbeddingsResponse(response: unknown): response is OpenAICreateEmbeddingsObject {
+  return (
+    response !== null &&
+    typeof response === 'object' &&
+    'object' in response &&
+    (response as Record<string, unknown>).object === 'list'
+  );
+}
+
 /**
  * Check if streaming event is from the Responses API
  */
