Skip to content

Commit

Permalink
feat(node): Add @vercel/ai instrumentation
Browse files Browse the repository at this point in the history
  • Loading branch information
AbhiPrasad committed Oct 7, 2024
1 parent 3f0926e commit 38cb4d4
Show file tree
Hide file tree
Showing 2 changed files with 69 additions and 0 deletions.
2 changes: 2 additions & 0 deletions packages/node/src/integrations/tracing/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ import { instrumentMysql2, mysql2Integration } from './mysql2';
import { instrumentNest, nestIntegration } from './nest/nest';
import { instrumentPostgres, postgresIntegration } from './postgres';
import { instrumentRedis, redisIntegration } from './redis';
import { vercelAiIntegration } from './vercelai';

/**
* With OTEL, all performance integrations will be added, as OTEL only initializes them when the patched package is actually required.
Expand Down Expand Up @@ -45,6 +46,7 @@ export function getAutoPerformanceIntegrations(): Integration[] {
kafkaIntegration(),
amqplibIntegration(),
lruMemoizerIntegration(),
vercelAiIntegration(),
];
}

Expand Down
67 changes: 67 additions & 0 deletions packages/node/src/integrations/tracing/vercelai.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,67 @@
import { SEMANTIC_ATTRIBUTE_SENTRY_OP, defineIntegration, spanToJSON } from '@sentry/core';
import type { IntegrationFn } from '@sentry/types';

import { addOriginToSpan } from '../../utils/addOriginToSpan';

const INTEGRATION_NAME = 'VercelAi';

const _vercelAiIntegration = (() => {
return {
name: INTEGRATION_NAME,
setupOnce() {
// TODO: Vercel Instrumentation.
},

setup(client) {
// TODO: Only run if @vercel/ai is installed.
client.on('spanStart', span => {
const { data: attributes, description: name } = spanToJSON(span);

if (!attributes || !name) {
return;
}

// https://sdk.vercel.ai/docs/ai-sdk-core/telemetry#basic-llm-span-information
if (attributes['ai.model.id'] && attributes['ai.model.provider']) {
addOriginToSpan(span, 'auto.vercel.ai');

const aiOperation = name.replace('ai.', '');
if (aiOperation.includes('embed')) {
span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.embeddings');
} else {
span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.run');
}

if (attributes['ai.prompt']) {
span.setAttribute('ai.input_messages', attributes['ai.prompt']);
}
if (attributes['ai.usage.completionTokens'] != undefined) {
span.setAttribute('ai.completion_tokens.used', attributes['ai.usage.completionTokens']);
}
if (attributes['ai.usage.promptTokens'] != undefined) {
span.setAttribute('ai.prompt_tokens.used', attributes['ai.usage.promptTokens']);
}
if (
attributes['ai.usage.completionTokens'] != undefined &&
attributes['ai.usage.promptTokens'] != undefined
) {
span.setAttribute(
'ai.tokens.used',
attributes['ai.usage.completionTokens'] + attributes['ai.usage.promptTokens'],
);
}
if (attributes['ai.model.id']) {
span.setAttribute('ai.model_id', attributes['ai.model.id']);
}

span.setAttribute('ai.streaming', name.includes('stream'));
}
});
},
};
}) satisfies IntegrationFn;

/**
* Integration for @vercel/ai
*/
export const vercelAiIntegration = defineIntegration(_vercelAiIntegration);

0 comments on commit 38cb4d4

Please sign in to comment.