UNPKG

@sentry/core

Version:
250 lines (222 loc) 9.09 kB
import { GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE, GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE, GEN_AI_INPUT_MESSAGES_ATTRIBUTE, GEN_AI_SYSTEM_INSTRUCTIONS_ATTRIBUTE, GEN_AI_INPUT_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE, GEN_AI_EXECUTE_TOOL_OPERATION_ATTRIBUTE, GEN_AI_RERANK_DO_RERANK_OPERATION_ATTRIBUTE, GEN_AI_EMBED_MANY_DO_EMBED_OPERATION_ATTRIBUTE, GEN_AI_EMBED_DO_EMBED_OPERATION_ATTRIBUTE, GEN_AI_STREAM_OBJECT_DO_STREAM_OPERATION_ATTRIBUTE, GEN_AI_GENERATE_OBJECT_DO_GENERATE_OPERATION_ATTRIBUTE, GEN_AI_STREAM_TEXT_DO_STREAM_OPERATION_ATTRIBUTE, GEN_AI_GENERATE_TEXT_DO_GENERATE_OPERATION_ATTRIBUTE, GEN_AI_INVOKE_AGENT_OPERATION_ATTRIBUTE } from '../ai/gen-ai-attributes.js'; import { extractSystemInstructions, getTruncatedJsonString } from '../ai/utils.js'; import { toolCallSpanContextMap } from './constants.js'; import { AI_PROMPT_ATTRIBUTE, AI_PROMPT_MESSAGES_ATTRIBUTE } from './vercel-ai-attributes.js'; /** * Accumulates token data from a span to its parent in the token accumulator map. * This function extracts token usage from the current span and adds it to the * accumulated totals for its parent span. */ function accumulateTokensForParent(span, tokenAccumulator) { const parentSpanId = span.parent_span_id; if (!parentSpanId) { return; } const inputTokens = span.data[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE]; const outputTokens = span.data[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE]; if (typeof inputTokens === 'number' || typeof outputTokens === 'number') { const existing = tokenAccumulator.get(parentSpanId) || { inputTokens: 0, outputTokens: 0 }; if (typeof inputTokens === 'number') { existing.inputTokens += inputTokens; } if (typeof outputTokens === 'number') { existing.outputTokens += outputTokens; } tokenAccumulator.set(parentSpanId, existing); } } /** * Applies accumulated token data to the `gen_ai.invoke_agent` span. * Only immediate children of the `gen_ai.invoke_agent` span are considered, * since aggregation will automatically occur for each parent span. */ function applyAccumulatedTokens( spanOrTrace, tokenAccumulator, ) { const accumulated = tokenAccumulator.get(spanOrTrace.span_id); if (!accumulated || !spanOrTrace.data) { return; } if (accumulated.inputTokens > 0) { spanOrTrace.data[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE] = accumulated.inputTokens; } if (accumulated.outputTokens > 0) { spanOrTrace.data[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE] = accumulated.outputTokens; } if (accumulated.inputTokens > 0 || accumulated.outputTokens > 0) { spanOrTrace.data['gen_ai.usage.total_tokens'] = accumulated.inputTokens + accumulated.outputTokens; } } /** * Get the span context associated with a tool call ID. */ function _INTERNAL_getSpanContextForToolCallId(toolCallId) { return toolCallSpanContextMap.get(toolCallId); } /** * Clean up the span mapping for a tool call ID */ function _INTERNAL_cleanupToolCallSpanContext(toolCallId) { toolCallSpanContextMap.delete(toolCallId); } /** * Convert an array of tool strings to a JSON string */ function convertAvailableToolsToJsonString(tools) { const toolObjects = tools.map(tool => { if (typeof tool === 'string') { try { return JSON.parse(tool); } catch { return tool; } } return tool; }); return JSON.stringify(toolObjects); } /** * Filter out invalid entries in messages array * @param input - The input array to filter * @returns The filtered array */ function filterMessagesArray(input) { return input.filter( (m) => !!m && typeof m === 'object' && 'role' in m && 'content' in m, ); } /** * Normalize the user input (stringified object with prompt, system, messages) to messages array */ function convertUserInputToMessagesFormat(userInput) { try { const p = JSON.parse(userInput); if (!!p && typeof p === 'object') { let { messages } = p; const { prompt, system } = p; const result = []; // prepend top-level system instruction if present if (typeof system === 'string') { result.push({ role: 'system', content: system }); } // stringified messages array if (typeof messages === 'string') { try { messages = JSON.parse(messages); } catch { // ignore parse errors } } // messages array format: { messages: [...] } if (Array.isArray(messages)) { result.push(...filterMessagesArray(messages)); return result; } // prompt array format: { prompt: [...] } if (Array.isArray(prompt)) { result.push(...filterMessagesArray(prompt)); return result; } // prompt string format: { prompt: "..." } if (typeof prompt === 'string') { result.push({ role: 'user', content: prompt }); } if (result.length > 0) { return result; } } // eslint-disable-next-line no-empty } catch {} return []; } /** * Generate a request.messages JSON array from the prompt field in the * invoke_agent op */ function requestMessagesFromPrompt(span, attributes) { if ( typeof attributes[AI_PROMPT_ATTRIBUTE] === 'string' && !attributes[GEN_AI_INPUT_MESSAGES_ATTRIBUTE] && !attributes[AI_PROMPT_MESSAGES_ATTRIBUTE] ) { // No messages array is present, so we need to convert the prompt to the proper messages format // This is the case for ai.generateText spans // The ai.prompt attribute is a stringified object with prompt, system, messages attributes // The format of these is described in the vercel docs, for instance: https://ai-sdk.dev/docs/reference/ai-sdk-core/stream-object#parameters const userInput = attributes[AI_PROMPT_ATTRIBUTE]; const messages = convertUserInputToMessagesFormat(userInput); if (messages.length) { const { systemInstructions, filteredMessages } = extractSystemInstructions(messages); if (systemInstructions) { span.setAttribute(GEN_AI_SYSTEM_INSTRUCTIONS_ATTRIBUTE, systemInstructions); } const filteredLength = Array.isArray(filteredMessages) ? filteredMessages.length : 0; const truncatedMessages = getTruncatedJsonString(filteredMessages); span.setAttributes({ [AI_PROMPT_ATTRIBUTE]: truncatedMessages, [GEN_AI_INPUT_MESSAGES_ATTRIBUTE]: truncatedMessages, [GEN_AI_INPUT_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE]: filteredLength, }); } } else if (typeof attributes[AI_PROMPT_MESSAGES_ATTRIBUTE] === 'string') { // In this case we already get a properly formatted messages array, this is the preferred way to get the messages // This is the case for ai.generateText.doGenerate spans try { const messages = JSON.parse(attributes[AI_PROMPT_MESSAGES_ATTRIBUTE]); if (Array.isArray(messages)) { const { systemInstructions, filteredMessages } = extractSystemInstructions(messages); if (systemInstructions) { span.setAttribute(GEN_AI_SYSTEM_INSTRUCTIONS_ATTRIBUTE, systemInstructions); } const filteredLength = Array.isArray(filteredMessages) ? filteredMessages.length : 0; const truncatedMessages = getTruncatedJsonString(filteredMessages); span.setAttributes({ [AI_PROMPT_MESSAGES_ATTRIBUTE]: truncatedMessages, [GEN_AI_INPUT_MESSAGES_ATTRIBUTE]: truncatedMessages, [GEN_AI_INPUT_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE]: filteredLength, }); } // eslint-disable-next-line no-empty } catch {} } } /** * Maps a Vercel AI span name to the corresponding Sentry op. */ function getSpanOpFromName(name) { switch (name) { case 'ai.generateText': case 'ai.streamText': case 'ai.generateObject': case 'ai.streamObject': case 'ai.embed': case 'ai.embedMany': case 'ai.rerank': return GEN_AI_INVOKE_AGENT_OPERATION_ATTRIBUTE; case 'ai.generateText.doGenerate': return GEN_AI_GENERATE_TEXT_DO_GENERATE_OPERATION_ATTRIBUTE; case 'ai.streamText.doStream': return GEN_AI_STREAM_TEXT_DO_STREAM_OPERATION_ATTRIBUTE; case 'ai.generateObject.doGenerate': return GEN_AI_GENERATE_OBJECT_DO_GENERATE_OPERATION_ATTRIBUTE; case 'ai.streamObject.doStream': return GEN_AI_STREAM_OBJECT_DO_STREAM_OPERATION_ATTRIBUTE; case 'ai.embed.doEmbed': return GEN_AI_EMBED_DO_EMBED_OPERATION_ATTRIBUTE; case 'ai.embedMany.doEmbed': return GEN_AI_EMBED_MANY_DO_EMBED_OPERATION_ATTRIBUTE; case 'ai.rerank.doRerank': return GEN_AI_RERANK_DO_RERANK_OPERATION_ATTRIBUTE; case 'ai.toolCall': return GEN_AI_EXECUTE_TOOL_OPERATION_ATTRIBUTE; default: if (name.startsWith('ai.stream')) { return 'ai.run'; } return undefined; } } export { _INTERNAL_cleanupToolCallSpanContext, _INTERNAL_getSpanContextForToolCallId, accumulateTokensForParent, applyAccumulatedTokens, convertAvailableToolsToJsonString, convertUserInputToMessagesFormat, getSpanOpFromName, requestMessagesFromPrompt }; //# sourceMappingURL=utils.js.map