Skip to content

Commit 71b6463

Browse files
nicohrubecnicohrubec
committed
attribute renames
1 parent d3901e3 commit 71b6463

5 files changed

Lines changed: 223 additions & 9 deletions

File tree

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
import * as Sentry from '@sentry/node';
2+
import { embed, embedMany } from 'ai';
3+
import { MockEmbeddingModelV1 } from 'ai/test';
4+
5+
async function run() {
6+
await Sentry.startSpan({ op: 'function', name: 'main' }, async () => {
7+
// Single embedding
8+
await embed({
9+
model: new MockEmbeddingModelV1({
10+
doEmbed: async () => ({
11+
embeddings: [[0.1, 0.2, 0.3]],
12+
usage: { tokens: 10 },
13+
}),
14+
}),
15+
value: 'Embedding test!',
16+
});
17+
18+
// Multiple embeddings
19+
await embedMany({
20+
model: new MockEmbeddingModelV1({
21+
maxEmbeddingsPerCall: 5,
22+
doEmbed: async () => ({
23+
embeddings: [
24+
[0.1, 0.2, 0.3],
25+
[0.4, 0.5, 0.6],
26+
],
27+
usage: { tokens: 20 },
28+
}),
29+
}),
30+
values: ['First input', 'Second input'],
31+
});
32+
});
33+
}
34+
35+
run();

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

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import { SEMANTIC_ATTRIBUTE_SENTRY_OP, SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN } from '
22
import type { Event } from '@sentry/node';
33
import { afterAll, describe, expect } from 'vitest';
44
import {
5+
GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE,
56
GEN_AI_INPUT_MESSAGES_ATTRIBUTE,
67
GEN_AI_INPUT_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE,
78
GEN_AI_OPERATION_NAME_ATTRIBUTE,
@@ -830,4 +831,140 @@ describe('Vercel AI integration', () => {
830831
});
831832
},
832833
);
834+
835+
createEsmAndCjsTests(__dirname, 'scenario-embeddings.mjs', 'instrument.mjs', (createRunner, test) => {
836+
test('creates embedding related spans with sendDefaultPii: false', async () => {
837+
const expectedTransaction = {
838+
transaction: 'main',
839+
spans: expect.arrayContaining([
840+
// embed invoke_agent span
841+
expect.objectContaining({
842+
data: expect.objectContaining({
843+
[GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'invoke_agent',
844+
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'gen_ai.invoke_agent',
845+
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.vercelai.otel',
846+
}),
847+
description: 'invoke_agent',
848+
op: 'gen_ai.invoke_agent',
849+
origin: 'auto.vercelai.otel',
850+
status: 'ok',
851+
}),
852+
// embed doEmbed span
853+
expect.objectContaining({
854+
data: expect.objectContaining({
855+
[GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'embeddings',
856+
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'gen_ai.embeddings',
857+
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.vercelai.otel',
858+
[GEN_AI_REQUEST_MODEL_ATTRIBUTE]: 'mock-model-id',
859+
[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE]: 10,
860+
[GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE]: 10,
861+
}),
862+
description: 'embeddings mock-model-id',
863+
op: 'gen_ai.embeddings',
864+
origin: 'auto.vercelai.otel',
865+
status: 'ok',
866+
}),
867+
// embedMany invoke_agent span
868+
expect.objectContaining({
869+
data: expect.objectContaining({
870+
[GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'invoke_agent',
871+
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'gen_ai.invoke_agent',
872+
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.vercelai.otel',
873+
}),
874+
description: 'invoke_agent',
875+
op: 'gen_ai.invoke_agent',
876+
origin: 'auto.vercelai.otel',
877+
status: 'ok',
878+
}),
879+
// embedMany doEmbed span
880+
expect.objectContaining({
881+
data: expect.objectContaining({
882+
[GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'embeddings',
883+
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'gen_ai.embeddings',
884+
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.vercelai.otel',
885+
[GEN_AI_REQUEST_MODEL_ATTRIBUTE]: 'mock-model-id',
886+
[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE]: 20,
887+
[GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE]: 20,
888+
}),
889+
description: 'embeddings mock-model-id',
890+
op: 'gen_ai.embeddings',
891+
origin: 'auto.vercelai.otel',
892+
status: 'ok',
893+
}),
894+
]),
895+
};
896+
897+
await createRunner().expect({ transaction: expectedTransaction }).start().completed();
898+
});
899+
});
900+
901+
createEsmAndCjsTests(__dirname, 'scenario-embeddings.mjs', 'instrument-with-pii.mjs', (createRunner, test) => {
902+
test('creates embedding related spans with sendDefaultPii: true', async () => {
903+
const expectedTransaction = {
904+
transaction: 'main',
905+
spans: expect.arrayContaining([
906+
// embed invoke_agent span with input
907+
expect.objectContaining({
908+
data: expect.objectContaining({
909+
[GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'invoke_agent',
910+
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'gen_ai.invoke_agent',
911+
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.vercelai.otel',
912+
[GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE]: 'Embedding test!',
913+
}),
914+
description: 'invoke_agent',
915+
op: 'gen_ai.invoke_agent',
916+
origin: 'auto.vercelai.otel',
917+
status: 'ok',
918+
}),
919+
// embed doEmbed span with input
920+
expect.objectContaining({
921+
data: expect.objectContaining({
922+
[GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'embeddings',
923+
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'gen_ai.embeddings',
924+
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.vercelai.otel',
925+
[GEN_AI_REQUEST_MODEL_ATTRIBUTE]: 'mock-model-id',
926+
[GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE]: '["Embedding test!"]',
927+
[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE]: 10,
928+
[GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE]: 10,
929+
}),
930+
description: 'embeddings mock-model-id',
931+
op: 'gen_ai.embeddings',
932+
origin: 'auto.vercelai.otel',
933+
status: 'ok',
934+
}),
935+
// embedMany invoke_agent span with input
936+
expect.objectContaining({
937+
data: expect.objectContaining({
938+
[GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'invoke_agent',
939+
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'gen_ai.invoke_agent',
940+
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.vercelai.otel',
941+
[GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE]: '["First input","Second input"]',
942+
}),
943+
description: 'invoke_agent',
944+
op: 'gen_ai.invoke_agent',
945+
origin: 'auto.vercelai.otel',
946+
status: 'ok',
947+
}),
948+
// embedMany doEmbed span with input
949+
expect.objectContaining({
950+
data: expect.objectContaining({
951+
[GEN_AI_OPERATION_NAME_ATTRIBUTE]: 'embeddings',
952+
[SEMANTIC_ATTRIBUTE_SENTRY_OP]: 'gen_ai.embeddings',
953+
[SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN]: 'auto.vercelai.otel',
954+
[GEN_AI_REQUEST_MODEL_ATTRIBUTE]: 'mock-model-id',
955+
[GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE]: '["First input","Second input"]',
956+
[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE]: 20,
957+
[GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE]: 20,
958+
}),
959+
description: 'embeddings mock-model-id',
960+
op: 'gen_ai.embeddings',
961+
origin: 'auto.vercelai.otel',
962+
status: 'ok',
963+
}),
964+
]),
965+
};
966+
967+
await createRunner().expect({ transaction: expectedTransaction }).start().completed();
968+
});
969+
});
833970
});

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -227,12 +227,12 @@ export const GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE = 'gen_ai.embeddings.input';
227227
/**
228228
* The span operation name for embedding
229229
*/
230-
export const GEN_AI_EMBED_DO_EMBED_OPERATION_ATTRIBUTE = 'gen_ai.embed';
230+
export const GEN_AI_EMBED_DO_EMBED_OPERATION_ATTRIBUTE = 'gen_ai.embeddings';
231231

232232
/**
233233
* The span operation name for embedding many
234234
*/
235-
export const GEN_AI_EMBED_MANY_DO_EMBED_OPERATION_ATTRIBUTE = 'gen_ai.embed_many';
235+
export const GEN_AI_EMBED_MANY_DO_EMBED_OPERATION_ATTRIBUTE = 'gen_ai.embeddings';
236236

237237
/**
238238
* The span operation name for reranking

‎packages/core/src/tracing/vercel-ai/constants.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ export const EMBEDDINGS_OPS = new Set(['ai.embed.doEmbed', 'ai.embedMany.doEmbed
2828
export const RERANK_OPS = new Set(['ai.rerank.doRerank']);
2929

3030
export const DO_SPAN_NAME_PREFIX: Record<string, string> = {
31-
'ai.embed.doEmbed': 'embed',
32-
'ai.embedMany.doEmbed': 'embed_many',
31+
'ai.embed.doEmbed': 'embeddings',
32+
'ai.embedMany.doEmbed': 'embeddings',
3333
'ai.rerank.doRerank': 'rerank',
3434
};

‎packages/core/src/tracing/vercel-ai/index.ts‎

Lines changed: 47 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import type { Event } from '../../types-hoist/event';
55
import type { Span, SpanAttributes, SpanAttributeValue, SpanJSON } from '../../types-hoist/span';
66
import { spanToJSON } from '../../utils/spanUtils';
77
import {
8+
GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE,
89
GEN_AI_INPUT_MESSAGES_ATTRIBUTE,
910
GEN_AI_OPERATION_NAME_ATTRIBUTE,
1011
GEN_AI_REQUEST_MODEL_ATTRIBUTE,
@@ -55,6 +56,9 @@ import {
5556
AI_USAGE_CACHED_INPUT_TOKENS_ATTRIBUTE,
5657
AI_USAGE_COMPLETION_TOKENS_ATTRIBUTE,
5758
AI_USAGE_PROMPT_TOKENS_ATTRIBUTE,
59+
AI_USAGE_TOKENS_ATTRIBUTE,
60+
AI_VALUE_ATTRIBUTE,
61+
AI_VALUES_ATTRIBUTE,
5862
OPERATION_NAME_ATTRIBUTE,
5963
} from './vercel-ai-attributes';
6064

@@ -160,6 +164,9 @@ function processEndedVercelAiSpan(span: SpanJSON): void {
160164
renameAttributeKey(attributes, 'ai.usage.inputTokens', GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE);
161165
renameAttributeKey(attributes, 'ai.usage.outputTokens', GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE);
162166

167+
// Embedding spans use ai.usage.tokens instead of promptTokens/completionTokens
168+
renameAttributeKey(attributes, AI_USAGE_TOKENS_ATTRIBUTE, GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE);
169+
163170
// AI SDK uses avgOutputTokensPerSecond, map to our expected attribute name
164171
renameAttributeKey(attributes, 'ai.response.avgOutputTokensPerSecond', 'ai.response.avgCompletionTokensPerSecond');
165172

@@ -172,12 +179,14 @@ function processEndedVercelAiSpan(span: SpanJSON): void {
172179
attributes[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE] + attributes[GEN_AI_USAGE_INPUT_TOKENS_CACHED_ATTRIBUTE];
173180
}
174181

175-
if (
176-
typeof attributes[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE] === 'number' &&
177-
typeof attributes[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE] === 'number'
178-
) {
182+
// Compute total tokens from input + output (embeddings may only have input tokens)
183+
if (typeof attributes[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE] === 'number') {
184+
const outputTokens =
185+
typeof attributes[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE] === 'number'
186+
? (attributes[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE] as number)
187+
: 0;
179188
attributes[GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE] =
180-
attributes[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE] + attributes[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE];
189+
outputTokens + (attributes[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE] as number);
181190
}
182191

183192
// Convert the available tools array to a JSON string
@@ -207,6 +216,39 @@ function processEndedVercelAiSpan(span: SpanJSON): void {
207216
renameAttributeKey(attributes, AI_SCHEMA_ATTRIBUTE, 'gen_ai.request.schema');
208217
renameAttributeKey(attributes, AI_MODEL_ID_ATTRIBUTE, GEN_AI_REQUEST_MODEL_ATTRIBUTE);
209218

219+
// Map embedding input: ai.value (single) or ai.values (batch) → gen_ai.embeddings.input
220+
// Vercel AI SDK JSON-stringifies each value individually, so for ai.value we parse it back,
221+
// and for ai.values we parse each element and re-stringify as a proper JSON array.
222+
if (attributes[AI_VALUE_ATTRIBUTE] != null) {
223+
const raw = attributes[AI_VALUE_ATTRIBUTE];
224+
if (typeof raw === 'string') {
225+
try {
226+
attributes[GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE] = JSON.parse(raw);
227+
} catch {
228+
attributes[GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE] = raw;
229+
}
230+
} else {
231+
attributes[GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE] = raw;
232+
}
233+
} else if (attributes[AI_VALUES_ATTRIBUTE] != null) {
234+
const values = attributes[AI_VALUES_ATTRIBUTE];
235+
if (Array.isArray(values)) {
236+
const parsed = values.map(v => {
237+
if (typeof v === 'string') {
238+
try {
239+
return JSON.parse(v);
240+
} catch {
241+
return v;
242+
}
243+
}
244+
return v;
245+
});
246+
attributes[GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE] = JSON.stringify(parsed);
247+
} else {
248+
attributes[GEN_AI_EMBEDDINGS_INPUT_ATTRIBUTE] = values;
249+
}
250+
}
251+
210252
addProviderMetadataToAttributes(attributes);
211253

212254
// Change attributes namespaced with `ai.X` to `vercel.ai.X`

0 commit comments

Comments
 (0)