Skip to content

Commit 935ef55

Browse files
nicohrubecnicohrubec
authored
feat(core): Support OpenAI embeddings API (#18224)
This adds instrumentation for the OpenAI Embeddings API. Specifically, we instrument [Create embeddings](https://platform.openai.com/docs/api-reference/embeddings/create), which is also the only endpoint in the embeddings API atm. Implementation generally follows the same flow we also have for the `completions` and `responses` APIs. To detect `embedding` requests we check whether the model name contains `embeddings`. The embedding results are currently not tracked, as we do not truncate outputs right now as far as I know and these can get large quite easily. For instance, [text-embedding-3 uses dimension 1536 (small) or 3072 (large) by default](https://platform.openai.com/docs/guides/embeddings#use-cases), resulting in single embeddings sizes of 6KB and 12KB, respectively. Test updates: - Added a new scenario-embeddings.mjs file, that covers the embeddings API tests (tried to put this in the main scenario.mjs, but the linter starts complaining about the file being too long). - Added a new scenario file to check that truncation works properly for the embeddings API. Also moved all truncation scenarios to a folder.
1 parent 610ae69 commit 935ef55

File tree

11 files changed

+432
-94
lines changed

11 files changed

+432
-94
lines changed

‎dev-packages/node-integration-tests/suites/tracing/openai/scenario.mjs‎ renamed to ‎dev-packages/node-integration-tests/suites/tracing/openai/scenario-chat.mjs‎

File renamed without changes.

‎dev-packages/node-integration-tests/suites/tracing/openai/scenario-embeddings.mjs‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
import { instrumentOpenAiClient } from '@sentry/core';
2+
import * as Sentry from '@sentry/node';
3+
4+
class MockOpenAI {
5+
constructor(config) {
6+
this.apiKey = config.apiKey;
7+
8+
this.embeddings = {
9+
create: async params => {
10+
await new Promise(resolve => setTimeout(resolve, 10));
11+
12+
if (params.model === 'error-model') {
13+
const error = new Error('Model not found');
14+
error.status = 404;
15+
error.headers = { 'x-request-id': 'mock-request-123' };
16+
throw error;
17+
}
18+
19+
return {
20+
object: 'list',
21+
data: [
22+
{
23+
object: 'embedding',
24+
embedding: [0.1, 0.2, 0.3],
25+
index: 0,
26+
},
27+
],
28+
model: params.model,
29+
usage: {
30+
prompt_tokens: 10,
31+
total_tokens: 10,
32+
},
33+
};
34+
},
35+
};
36+
}
37+
}
38+
39+
async function run() {
40+
await Sentry.startSpan({ op: 'function', name: 'main' }, async () => {
41+
const mockClient = new MockOpenAI({
42+
apiKey: 'mock-api-key',
43+
});
44+
45+
const client = instrumentOpenAiClient(mockClient);
46+
47+
// First test: embeddings API
48+
await client.embeddings.create({
49+
input: 'Embedding test!',
50+
model: 'text-embedding-3-small',
51+
dimensions: 1536,
52+
encoding_format: 'float',
53+
});
54+
55+
// Second test: embeddings API error model
56+
try {
57+
await client.embeddings.create({
58+
input: 'Error embedding test!',
59+
model: 'error-model',
60+
});
61+
} catch {
62+
// Error is expected and handled
63+
}
64+
});
65+
}
66+
67+
run();

‎dev-packages/node-integration-tests/suites/tracing/openai/test.ts‎

Lines changed: 138 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ describe('OpenAI integration', () => {
66
cleanupChildProcesses();
77
});
88

9-
const EXPECTED_TRANSACTION_DEFAULT_PII_FALSE = {
9+
const EXPECTED_TRANSACTION_DEFAULT_PII_FALSE_CHAT = {
1010
transaction: 'main',
1111
spans: expect.arrayContaining([
1212
// First span - basic chat completion without PII
@@ -147,7 +147,7 @@ describe('OpenAI integration', () => {
147147
]),
148148
};
149149

150-
const EXPECTED_TRANSACTION_DEFAULT_PII_TRUE = {
150+
const EXPECTED_TRANSACTION_DEFAULT_PII_TRUE_CHAT = {
151151
transaction: 'main',
152152
spans: expect.arrayContaining([
153153
// First span - basic chat completion with PII
@@ -321,27 +321,27 @@ describe('OpenAI integration', () => {
321321
]),
322322
};
323323

324-
createEsmAndCjsTests(__dirname, 'scenario.mjs', 'instrument.mjs', (createRunner, test) => {
324+
createEsmAndCjsTests(__dirname, 'scenario-chat.mjs', 'instrument.mjs', (createRunner, test) => {
325325
test('creates openai related spans with sendDefaultPii: false', async () => {
326326
await createRunner()
327327
.ignore('event')
328-
.expect({ transaction: EXPECTED_TRANSACTION_DEFAULT_PII_FALSE })
328+
.expect({ transaction: EXPECTED_TRANSACTION_DEFAULT_PII_FALSE_CHAT })
329329
.start()
330330
.completed();
331331
});
332332
});
333333

334-
createEsmAndCjsTests(__dirname, 'scenario.mjs', 'instrument-with-pii.mjs', (createRunner, test) => {
334+
createEsmAndCjsTests(__dirname, 'scenario-chat.mjs', 'instrument-with-pii.mjs', (createRunner, test) => {
335335
test('creates openai related spans with sendDefaultPii: true', async () => {
336336
await createRunner()
337337
.ignore('event')
338-
.expect({ transaction: EXPECTED_TRANSACTION_DEFAULT_PII_TRUE })
338+
.expect({ transaction: EXPECTED_TRANSACTION_DEFAULT_PII_TRUE_CHAT })
339339
.start()
340340
.completed();
341341
});
342342
});
343343

344-
createEsmAndCjsTests(__dirname, 'scenario.mjs', 'instrument-with-options.mjs', (createRunner, test) => {
344+
createEsmAndCjsTests(__dirname, 'scenario-chat.mjs', 'instrument-with-options.mjs', (createRunner, test) => {
345345
test('creates openai related spans with custom options', async () => {
346346
await createRunner()
347347
.ignore('event')
@@ -351,6 +351,109 @@ describe('OpenAI integration', () => {
351351
});
352352
});
353353

354+
const EXPECTED_TRANSACTION_DEFAULT_PII_FALSE_EMBEDDINGS = {
355+
transaction: 'main',
356+
spans: expect.arrayContaining([
357+
// First span - embeddings API
358+
expect.objectContaining({
359+
data: {
360+
'gen_ai.operation.name': 'embeddings',
361+
'sentry.op': 'gen_ai.embeddings',
362+
'sentry.origin': 'auto.ai.openai',
363+
'gen_ai.system': 'openai',
364+
'gen_ai.request.model': 'text-embedding-3-small',
365+
'gen_ai.request.encoding_format': 'float',
366+
'gen_ai.request.dimensions': 1536,
367+
'gen_ai.response.model': 'text-embedding-3-small',
368+
'gen_ai.usage.input_tokens': 10,
369+
'gen_ai.usage.total_tokens': 10,
370+
'openai.response.model': 'text-embedding-3-small',
371+
'openai.usage.prompt_tokens': 10,
372+
},
373+
description: 'embeddings text-embedding-3-small',
374+
op: 'gen_ai.embeddings',
375+
origin: 'auto.ai.openai',
376+
status: 'ok',
377+
}),
378+
// Second span - embeddings API error model
379+
expect.objectContaining({
380+
data: {
381+
'gen_ai.operation.name': 'embeddings',
382+
'sentry.op': 'gen_ai.embeddings',
383+
'sentry.origin': 'auto.ai.openai',
384+
'gen_ai.system': 'openai',
385+
'gen_ai.request.model': 'error-model',
386+
},
387+
description: 'embeddings error-model',
388+
op: 'gen_ai.embeddings',
389+
origin: 'auto.ai.openai',
390+
status: 'internal_error',
391+
}),
392+
]),
393+
};
394+
395+
const EXPECTED_TRANSACTION_DEFAULT_PII_TRUE_EMBEDDINGS = {
396+
transaction: 'main',
397+
spans: expect.arrayContaining([
398+
// First span - embeddings API with PII
399+
expect.objectContaining({
400+
data: {
401+
'gen_ai.operation.name': 'embeddings',
402+
'sentry.op': 'gen_ai.embeddings',
403+
'sentry.origin': 'auto.ai.openai',
404+
'gen_ai.system': 'openai',
405+
'gen_ai.request.model': 'text-embedding-3-small',
406+
'gen_ai.request.encoding_format': 'float',
407+
'gen_ai.request.dimensions': 1536,
408+
'gen_ai.request.messages': 'Embedding test!',
409+
'gen_ai.response.model': 'text-embedding-3-small',
410+
'gen_ai.usage.input_tokens': 10,
411+
'gen_ai.usage.total_tokens': 10,
412+
'openai.response.model': 'text-embedding-3-small',
413+
'openai.usage.prompt_tokens': 10,
414+
},
415+
description: 'embeddings text-embedding-3-small',
416+
op: 'gen_ai.embeddings',
417+
origin: 'auto.ai.openai',
418+
status: 'ok',
419+
}),
420+
// Second span - embeddings API error model with PII
421+
expect.objectContaining({
422+
data: {
423+
'gen_ai.operation.name': 'embeddings',
424+
'sentry.op': 'gen_ai.embeddings',
425+
'sentry.origin': 'auto.ai.openai',
426+
'gen_ai.system': 'openai',
427+
'gen_ai.request.model': 'error-model',
428+
'gen_ai.request.messages': 'Error embedding test!',
429+
},
430+
description: 'embeddings error-model',
431+
op: 'gen_ai.embeddings',
432+
origin: 'auto.ai.openai',
433+
status: 'internal_error',
434+
}),
435+
]),
436+
};
437+
createEsmAndCjsTests(__dirname, 'scenario-embeddings.mjs', 'instrument.mjs', (createRunner, test) => {
438+
test('creates openai related spans with sendDefaultPii: false', async () => {
439+
await createRunner()
440+
.ignore('event')
441+
.expect({ transaction: EXPECTED_TRANSACTION_DEFAULT_PII_FALSE_EMBEDDINGS })
442+
.start()
443+
.completed();
444+
});
445+
});
446+
447+
createEsmAndCjsTests(__dirname, 'scenario-embeddings.mjs', 'instrument-with-pii.mjs', (createRunner, test) => {
448+
test('creates openai related spans with sendDefaultPii: true', async () => {
449+
await createRunner()
450+
.ignore('event')
451+
.expect({ transaction: EXPECTED_TRANSACTION_DEFAULT_PII_TRUE_EMBEDDINGS })
452+
.start()
453+
.completed();
454+
});
455+
});
456+
354457
createEsmAndCjsTests(__dirname, 'scenario-root-span.mjs', 'instrument.mjs', (createRunner, test) => {
355458
test('it works without a wrapping span', async () => {
356459
await createRunner()
@@ -400,7 +503,7 @@ describe('OpenAI integration', () => {
400503

401504
createEsmAndCjsTests(
402505
__dirname,
403-
'scenario-message-truncation-completions.mjs',
506+
'truncation/scenario-message-truncation-completions.mjs',
404507
'instrument-with-pii.mjs',
405508
(createRunner, test) => {
406509
test('truncates messages when they exceed byte limit - keeps only last message and crops it', async () => {
@@ -436,7 +539,7 @@ describe('OpenAI integration', () => {
436539

437540
createEsmAndCjsTests(
438541
__dirname,
439-
'scenario-message-truncation-responses.mjs',
542+
'truncation/scenario-message-truncation-responses.mjs',
440543
'instrument-with-pii.mjs',
441544
(createRunner, test) => {
442545
test('truncates string inputs when they exceed byte limit', async () => {
@@ -469,4 +572,30 @@ describe('OpenAI integration', () => {
469572
});
470573
},
471574
);
575+
576+
createEsmAndCjsTests(
577+
__dirname,
578+
'truncation/scenario-message-truncation-embeddings.mjs',
579+
'instrument-with-pii.mjs',
580+
(createRunner, test) => {
581+
test('truncates messages when they exceed byte limit - keeps only last message and crops it', async () => {
582+
await createRunner()
583+
.ignore('event')
584+
.expect({
585+
transaction: {
586+
transaction: 'main',
587+
spans: expect.arrayContaining([
588+
expect.objectContaining({
589+
data: expect.objectContaining({
590+
'gen_ai.operation.name': 'embeddings',
591+
}),
592+
}),
593+
]),
594+
},
595+
})
596+
.start()
597+
.completed();
598+
});
599+
},
600+
);
472601
});

‎dev-packages/node-integration-tests/suites/tracing/openai/scenario-message-truncation-completions.mjs‎ renamed to ‎dev-packages/node-integration-tests/suites/tracing/openai/truncation/scenario-message-truncation-completions.mjs‎

File renamed without changes.

‎dev-packages/node-integration-tests/suites/tracing/openai/truncation/scenario-message-truncation-embeddings.mjs‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
import { instrumentOpenAiClient } from '@sentry/core';
2+
import * as Sentry from '@sentry/node';
3+
4+
class MockOpenAI {
5+
constructor(config) {
6+
this.apiKey = config.apiKey;
7+
8+
this.embeddings = {
9+
create: async params => {
10+
await new Promise(resolve => setTimeout(resolve, 10));
11+
12+
return {
13+
object: 'list',
14+
data: [
15+
{
16+
object: 'embedding',
17+
embedding: [0.1, 0.2, 0.3],
18+
index: 0,
19+
},
20+
],
21+
model: params.model,
22+
usage: {
23+
prompt_tokens: 10,
24+
total_tokens: 10,
25+
},
26+
};
27+
},
28+
};
29+
}
30+
}
31+
32+
async function run() {
33+
await Sentry.startSpan({ op: 'function', name: 'main' }, async () => {
34+
const mockClient = new MockOpenAI({
35+
apiKey: 'mock-api-key',
36+
});
37+
38+
const client = instrumentOpenAiClient(mockClient);
39+
40+
// Create 1 large input that gets truncated to fit within the 20KB limit
41+
const largeContent = 'A'.repeat(25000) + 'B'.repeat(25000); // ~50KB gets truncated to include only As
42+
43+
await client.embeddings.create({
44+
input: largeContent,
45+
model: 'text-embedding-3-small',
46+
dimensions: 1536,
47+
encoding_format: 'float',
48+
});
49+
50+
// Create 3 large inputs where:
51+
// - First 2 inputs are very large (will be dropped)
52+
// - Last input is large but will be truncated to fit within the 20KB limit
53+
const largeContent1 = 'A'.repeat(15000); // ~15KB
54+
const largeContent2 = 'B'.repeat(15000); // ~15KB
55+
const largeContent3 = 'C'.repeat(25000); // ~25KB (will be truncated)
56+
57+
await client.embeddings.create({
58+
input: [largeContent1, largeContent2, largeContent3],
59+
model: 'text-embedding-3-small',
60+
dimensions: 1536,
61+
encoding_format: 'float',
62+
});
63+
});
64+
}
65+
66+
run();

‎dev-packages/node-integration-tests/suites/tracing/openai/scenario-message-truncation-responses.mjs‎ renamed to ‎dev-packages/node-integration-tests/suites/tracing/openai/truncation/scenario-message-truncation-responses.mjs‎

File renamed without changes.

‎packages/core/src/tracing/ai/gen-ai-attributes.ts‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,16 @@ export const GEN_AI_REQUEST_TOP_K_ATTRIBUTE = 'gen_ai.request.top_k';
6565
*/
6666
export const GEN_AI_REQUEST_STOP_SEQUENCES_ATTRIBUTE = 'gen_ai.request.stop_sequences';
6767

68+
/**
69+
* The encoding format for the model request
70+
*/
71+
export const GEN_AI_REQUEST_ENCODING_FORMAT_ATTRIBUTE = 'gen_ai.request.encoding_format';
72+
73+
/**
74+
* The dimensions for the model request
75+
*/
76+
export const GEN_AI_REQUEST_DIMENSIONS_ATTRIBUTE = 'gen_ai.request.dimensions';
77+
6878
/**
6979
* Array of reasons why the model stopped generating tokens
7080
*/
@@ -208,6 +218,7 @@ export const OPENAI_USAGE_PROMPT_TOKENS_ATTRIBUTE = 'openai.usage.prompt_tokens'
208218
export const OPENAI_OPERATIONS = {
209219
CHAT: 'chat',
210220
RESPONSES: 'responses',
221+
EMBEDDINGS: 'embeddings',
211222
} as const;
212223

213224
// =============================================================================

‎packages/core/src/tracing/openai/constants.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ export const OPENAI_INTEGRATION_NAME = 'OpenAI';
22

33
// https://platform.openai.com/docs/quickstart?api-mode=responses
44
// https://platform.openai.com/docs/quickstart?api-mode=chat
5-
export const INSTRUMENTED_METHODS = ['responses.create', 'chat.completions.create'] as const;
5+
export const INSTRUMENTED_METHODS = ['responses.create', 'chat.completions.create', 'embeddings.create'] as const;
66
export const RESPONSES_TOOL_CALL_EVENT_TYPES = [
77
'response.output_item.added',
88
'response.function_call_arguments.delta',

0 commit comments

Comments
 (0)