diff --git a/packages/server-utils/src/integrations/aws-sdk/services/bedrock-runtime.ts b/packages/server-utils/src/integrations/aws-sdk/services/bedrock-runtime.ts index 7259dc085004..97854b017c8c 100644 --- a/packages/server-utils/src/integrations/aws-sdk/services/bedrock-runtime.ts +++ b/packages/server-utils/src/integrations/aws-sdk/services/bedrock-runtime.ts @@ -542,28 +542,36 @@ function recordLlamaAttributes(parsedChunk: ParsedChunk, span: Span): void { } } -function recordMistralAttributes(parsedChunk: ParsedChunk, span: Span): void { - if (parsedChunk.outputs?.[0]?.text !== undefined) { - span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(parsedChunk.outputs[0].text.length / 6)); +/** + * Bedrock adds the token counts to the last chunk of a stream. For models that send no counts of their own, + * this is the only real count. + */ +function recordInvocationMetrics(parsedChunk: ParsedChunk, span: Span): void { + const metrics = parsedChunk['amazon-bedrock-invocationMetrics']; + if (typeof metrics?.inputTokenCount === 'number') { + span.setAttribute(GEN_AI_USAGE_INPUT_TOKENS, metrics.inputTokenCount); } + if (typeof metrics?.outputTokenCount === 'number') { + span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, metrics.outputTokenCount); + } +} + +function recordMistralAttributes(parsedChunk: ParsedChunk, span: Span): void { + recordInvocationMetrics(parsedChunk, span); if (parsedChunk.outputs?.[0]?.stop_reason !== undefined) { span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [parsedChunk.outputs[0].stop_reason]); } } function recordCohereAttributes(parsedChunk: ParsedChunk, span: Span): void { - if (parsedChunk.generations?.[0]?.text !== undefined) { - span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(parsedChunk.generations[0].text.length / 6)); - } + recordInvocationMetrics(parsedChunk, span); if (parsedChunk.generations?.[0]?.finish_reason !== undefined) { span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [parsedChunk.generations[0].finish_reason]); } } function recordCohereRAttributes(parsedChunk: ParsedChunk, span: Span): void { - if (parsedChunk.text !== undefined) { - span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(parsedChunk.text.length / 6)); - } + recordInvocationMetrics(parsedChunk, span); if (parsedChunk.finish_reason !== undefined) { span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, [parsedChunk.finish_reason]); } diff --git a/packages/server-utils/test/integrations/aws-sdk-bedrock.test.ts b/packages/server-utils/test/integrations/aws-sdk-bedrock.test.ts index a1ba675251e6..a63e70d3324c 100644 --- a/packages/server-utils/test/integrations/aws-sdk-bedrock.test.ts +++ b/packages/server-utils/test/integrations/aws-sdk-bedrock.test.ts @@ -172,6 +172,54 @@ describe('BedrockRuntimeServiceExtension stream spans', () => { expect(spanToJSON(span).attributes?.[GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(42); }); + const INVOCATION_METRICS = { + inputTokenCount: 11, + outputTokenCount: 3, + invocationLatency: 412, + firstByteLatency: 198, + }; + + it.each( + [ + { modelId: 'mistral.mistral-large-2402-v1:0', toChunk: (text: string) => ({ outputs: [{ text }] }) }, + { modelId: 'cohere.command-text-v14', toChunk: (text: string) => ({ generations: [{ text }] }) }, + { modelId: 'cohere.command-r-plus-v1:0', toChunk: (text: string) => ({ text }) }, + ].flatMap(family => [ + { + ...family, + when: 'has invocation metrics', + metrics: INVOCATION_METRICS, + expected: { [GEN_AI_USAGE_INPUT_TOKENS]: 11, [GEN_AI_USAGE_OUTPUT_TOKENS]: 3 }, + }, + { + ...family, + when: 'has no invocation metrics', + metrics: undefined, + expected: { [GEN_AI_USAGE_INPUT_TOKENS]: undefined, [GEN_AI_USAGE_OUTPUT_TOKENS]: undefined }, + }, + ]), + )('records only real token counts when a $modelId stream $when', async ({ modelId, toChunk, metrics, expected }) => { + const span = startInactiveSpan({ name: 'bedrock' }); + const encoder = new TextEncoder(); + const chunks = [toChunk('The capital'), { ...toChunk(' is Paris.'), 'amazon-bedrock-invocationMetrics': metrics }]; + const data = { + body: (async function* () { + for (const chunk of chunks) { + yield { chunk: { bytes: encoder.encode(JSON.stringify(chunk)) } }; + } + })(), + }; + + respond(span, 'InvokeModelWithResponseStream', { modelId }, data); + await drain(data.body); + + const attributes = spanToJSON(span).attributes; + expect({ + [GEN_AI_USAGE_INPUT_TOKENS]: attributes?.[GEN_AI_USAGE_INPUT_TOKENS], + [GEN_AI_USAGE_OUTPUT_TOKENS]: attributes?.[GEN_AI_USAGE_OUTPUT_TOKENS], + }).toEqual(expected); + }); + it('ends an InvokeModelWithResponseStream span without a modelId when the stream is drained', async () => { const span = startInactiveSpan({ name: 'bedrock' }); const data = { body: (async function* () {})() };