From 315c3c0b5d3885b7e066d253f1118a94c74f804f Mon Sep 17 00:00:00 2001 From: Andrei Borza Date: Fri, 13 Nov 2026 15:47:08 +0100 Subject: [PATCH] feat(server-utils): Record Bedrock responses as output messages Co-Authored-By: Claude Opus 5.5 --- .../suites/aws-serverless/bedrock/test.ts | 13 +- .../services/bedrock-runtime-content.ts | 74 ++++- .../aws-sdk/services/bedrock-runtime.ts | 125 +++++--- .../test/integrations/aws-sdk-bedrock.test.ts | 273 ++++++++++++++++++ 4 files changed, 443 insertions(+), 42 deletions(-) diff --git a/dev-packages/node-integration-tests/suites/aws-serverless/bedrock/test.ts b/dev-packages/node-integration-tests/suites/aws-serverless/bedrock/test.ts index dd9867391034..937cbc4fe54d 100644 --- a/dev-packages/node-integration-tests/suites/aws-serverless/bedrock/test.ts +++ b/dev-packages/node-integration-tests/suites/aws-serverless/bedrock/test.ts @@ -1,6 +1,7 @@ import { GEN_AI_INPUT_MESSAGES, GEN_AI_OPERATION_NAME, + GEN_AI_OUTPUT_MESSAGES, GEN_AI_PROVIDER_NAME, GEN_AI_REQUEST_MAX_TOKENS, GEN_AI_REQUEST_MODEL, @@ -59,7 +60,11 @@ function assertBedrockSpans(container: SerializedStreamedSpanContainer): void { const SYSTEM_INSTRUCTIONS = JSON.stringify([{ type: 'text', content: 'You are a weather bot.' }]); -function assertBedrockInputs(container: SerializedStreamedSpanContainer): void { +const OUTPUT_MESSAGES = JSON.stringify([ + { role: 'assistant', parts: [{ type: 'text', content: 'Hello from Bedrock' }], finish_reason: 'end_turn' }, +]); + +function assertBedrockContent(container: SerializedStreamedSpanContainer): void { const converseSpan = container.items.find(span => span.name === `chat ${MODEL_ID}`); expect(converseSpan!.attributes[GEN_AI_SYSTEM_INSTRUCTIONS].value).toBe(SYSTEM_INSTRUCTIONS); expect(converseSpan!.attributes[GEN_AI_INPUT_MESSAGES].value).toBe( @@ -68,12 +73,14 @@ function assertBedrockInputs(container: SerializedStreamedSpanContainer): void { expect(converseSpan!.attributes[GEN_AI_TOOL_DEFINITIONS].value).toBe( JSON.stringify([{ name: 'get_weather', inputSchema: { json: { type: 'object' } } }]), ); + expect(converseSpan!.attributes[GEN_AI_OUTPUT_MESSAGES].value).toBe(OUTPUT_MESSAGES); const invokeModelSpan = container.items.find(span => span.name === `generate_content ${MODEL_ID}`); expect(invokeModelSpan!.attributes[GEN_AI_SYSTEM_INSTRUCTIONS].value).toBe(SYSTEM_INSTRUCTIONS); expect(invokeModelSpan!.attributes[GEN_AI_INPUT_MESSAGES].value).toBe( JSON.stringify([{ role: 'user', content: 'Hello' }]), ); + expect(invokeModelSpan!.attributes[GEN_AI_OUTPUT_MESSAGES].value).toBe(OUTPUT_MESSAGES); } function assertStreamSpans(container: SerializedStreamedSpanContainer): void { @@ -113,8 +120,8 @@ describe('awsIntegration - Bedrock', () => { }); createEsmAndCjsTests(__dirname, 'scenario.mjs', 'instrument-gen-ai.mjs', (createTestRunner, test) => { - test('records Bedrock inputs by default', { timeout: 90_000 }, async () => { - await createTestRunner().ignore('event').expect({ span: assertBedrockInputs }).start().completed(); + test('records Bedrock inputs and outputs by default', { timeout: 90_000 }, async () => { + await createTestRunner().ignore('event').expect({ span: assertBedrockContent }).start().completed(); }); }); diff --git a/packages/server-utils/src/integrations/aws-sdk/services/bedrock-runtime-content.ts b/packages/server-utils/src/integrations/aws-sdk/services/bedrock-runtime-content.ts index b38068553f21..d96527b21a44 100644 --- a/packages/server-utils/src/integrations/aws-sdk/services/bedrock-runtime-content.ts +++ b/packages/server-utils/src/integrations/aws-sdk/services/bedrock-runtime-content.ts @@ -1,6 +1,8 @@ +import type { Span } from '@sentry/core'; import { isObjectLike, stringify } from '@sentry/core'; import { GEN_AI_INPUT_MESSAGES, + GEN_AI_OUTPUT_MESSAGES, GEN_AI_SYSTEM_INSTRUCTIONS, GEN_AI_TOOL_DEFINITIONS, } from '@sentry/conventions/attributes'; @@ -11,11 +13,11 @@ import type { GenAiMessagePart } from '../../../ai/types'; * Translate Converse `ContentBlock`s (also used by Amazon Nova `InvokeModel` bodies) into message parts. * The conventions require a `type` on each part, and Bedrock's blocks (`{ text }`, `{ toolUse }`, ...) have none. * - * Media is reported by modality only: inline media is raw bytes, which the conventions require to be - * dropped rather than recorded. Blocks without content worth recording (`cachePoint`, redacted + * Media is reported by modality only: inline media is raw bytes, which would serialize as a large array of + * numbers. Blocks without content worth recording (`cachePoint`, redacted * reasoning, unknown members) are left out rather than serialized, so no binary payload slips through. */ -function contentBlocksToParts(content: unknown): GenAiMessagePart[] { +export function contentBlocksToParts(content: unknown): GenAiMessagePart[] { return Array.isArray(content) ? content.flatMap(contentBlockToParts) : []; } @@ -156,3 +158,69 @@ export function getPromptInputAttributes(prompt: unknown): Record { + if (!isObjectLike(block)) { + return []; + } + if (block.type === 'text') { + return textParts(block.text); + } + if (block.type === 'thinking') { + return typeof block.thinking === 'string' ? [{ type: 'reasoning', content: block.thinking }] : []; + } + if (block.type === 'tool_use') { + return [{ type: 'tool_call', id: block.id, name: block.name, arguments: stringify(block.input ?? {}, String) }]; + } + return []; + }); +} + +export function textParts(text: unknown): GenAiMessagePart[] { + return typeof text === 'string' ? [{ type: 'text', content: text }] : []; +} + +/** Record the parts as one assistant message. Parts of a stream are kept at their block index, so skip the gaps. */ +export function setOutputMessages(span: Span, parts: GenAiMessagePart[], finishReason: unknown): void { + const filledParts = parts.filter(Boolean); + if (filledParts.length) { + // An `undefined` finish reason is left out of the JSON. + span.setAttribute( + GEN_AI_OUTPUT_MESSAGES, + JSON.stringify([{ role: 'assistant', parts: filledParts, finish_reason: finishReason }]), + ); + } +} + +/** Append a streamed fragment to the part at the block `index`, and open a `type` part there if there is none. */ +export function appendStreamPart( + parts: GenAiMessagePart[] | undefined, + index: unknown, + type: 'text' | 'reasoning' | 'tool_call', + fragment: unknown, +): void { + if (!parts || typeof fragment !== 'string') { + return; + } + // Tool call arguments stream in as fragments of their JSON. + const field = type === 'tool_call' ? 'arguments' : 'content'; + const part = (parts[typeof index === 'number' ? index : 0] ??= { type, [field]: '' }); + part[field] = `${part[field]}${fragment}`; +} + +export function startStreamToolCall( + parts: GenAiMessagePart[] | undefined, + index: unknown, + id: unknown, + name: unknown, +): void { + if (parts) { + parts[typeof index === 'number' ? index : 0] = { type: 'tool_call', id, name, arguments: '' }; + } +} 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 f77e04f6f5ed..8345046d0146 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 @@ -14,6 +14,7 @@ import { } from '@sentry/conventions/attributes'; import { GEN_AI_CHAT, GEN_AI_GENERATE_CONTENT } from '@sentry/conventions/op'; import type { StreamResponseState } from '../../../ai/core/utils'; +import type { GenAiMessagePart } from '../../../ai/types'; import { endStreamSpan, getTokenUsageAttributes, resolveAIRecordingOptions } from '../../../ai/core/utils'; import { isInsidePiAiSpan } from '../../../ai/pi-ai/providers'; import { DEBUG_BUILD } from '../../../debug-build'; @@ -24,9 +25,15 @@ import { } from '../constants'; import type { NormalizedRequest, NormalizedResponse } from '../types'; import { + appendStreamPart, + contentBlocksToParts, getClaudeInputAttributes, + getClaudeOutputParts, getConverseInputAttributes, getPromptInputAttributes, + setOutputMessages, + startStreamToolCall, + textParts, } from './bedrock-runtime-content'; import type { RequestMetadata, ServiceExtension } from './ServiceExtension'; @@ -51,6 +58,9 @@ interface ConverseStreamOutput { // of one structural type per family. type ParsedChunk = any; +/** The output parts are only collected when outputs are recorded. */ +type BedrockStreamState = StreamResponseState & { parts?: GenAiMessagePart[] }; + const textDecoder = new TextDecoder(); export class BedrockRuntimeServiceExtension implements ServiceExtension { @@ -282,9 +292,12 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { } private _responseHookConverse(response: NormalizedResponse, span: Span): void { - const { stopReason, usage } = response.data; + const { stopReason, usage, output } = response.data; setStopReason(span, stopReason); + if (resolveAIRecordingOptions().recordOutputs) { + setOutputMessages(span, contentBlocksToParts(output?.message?.content), stopReason); + } span.setAttributes( getTokenUsageAttributes( // With prompt caching, `inputTokens` leaves out cache reads and writes, which the conventions count as input. @@ -319,6 +332,15 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { if (response.data?.body) { const decodedResponseBody = textDecoder.decode(response.data.body); const responseBody = JSON.parse(decodedResponseBody); + // Nova; Claude and Llama; Cohere Command R; Titan; Cohere Command; Mistral. + const finishReason = + responseBody.stopReason ?? + responseBody.stop_reason ?? + responseBody.finish_reason ?? + responseBody.results?.[0]?.completionReason ?? + responseBody.generations?.[0]?.finish_reason ?? + responseBody.outputs?.[0]?.stop_reason; + setStopReason(span, finishReason); if (currentModelId.includes('amazon.titan')) { if (responseBody.inputTextTokenCount !== undefined) { span.setAttribute(GEN_AI_USAGE_INPUT_TOKENS, responseBody.inputTextTokenCount); @@ -326,9 +348,6 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { if (responseBody.results?.[0]?.tokenCount !== undefined) { span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, responseBody.results[0].tokenCount); } - if (responseBody.results?.[0]?.completionReason !== undefined) { - span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, JSON.stringify([responseBody.results[0].completionReason])); - } } else if (currentModelId.includes('amazon.nova')) { if (responseBody.usage !== undefined) { if (responseBody.usage.inputTokens !== undefined) { @@ -338,9 +357,6 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, responseBody.usage.outputTokens); } } - if (responseBody.stopReason !== undefined) { - span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, JSON.stringify([responseBody.stopReason])); - } } else if (currentModelId.includes('anthropic.claude')) { const usage = responseBody.usage; span.setAttributes( @@ -354,9 +370,6 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { usage?.cache_read_input_tokens, ), ); - if (responseBody.stop_reason !== undefined) { - span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, JSON.stringify([responseBody.stop_reason])); - } } else if (currentModelId.includes('meta.llama')) { if (responseBody.prompt_token_count !== undefined) { span.setAttribute(GEN_AI_USAGE_INPUT_TOKENS, responseBody.prompt_token_count); @@ -364,33 +377,21 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { if (responseBody.generation_token_count !== undefined) { span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, responseBody.generation_token_count); } - if (responseBody.stop_reason !== undefined) { - span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, JSON.stringify([responseBody.stop_reason])); - } } else if (currentModelId.includes('cohere.command-r')) { if (responseBody.text !== undefined) { span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(responseBody.text.length / 6)); } - if (responseBody.finish_reason !== undefined) { - span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, JSON.stringify([responseBody.finish_reason])); - } } else if (currentModelId.includes('cohere.command')) { if (responseBody.generations?.[0]?.text !== undefined) { span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(responseBody.generations[0].text.length / 6)); } - if (responseBody.generations?.[0]?.finish_reason !== undefined) { - span.setAttribute( - GEN_AI_RESPONSE_FINISH_REASONS, - JSON.stringify([responseBody.generations[0].finish_reason]), - ); - } } else if (currentModelId.includes('mistral')) { if (responseBody.outputs?.[0]?.text !== undefined) { span.setAttribute(GEN_AI_USAGE_OUTPUT_TOKENS, Math.ceil(responseBody.outputs[0].text.length / 6)); } - if (responseBody.outputs?.[0]?.stop_reason !== undefined) { - span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, JSON.stringify([responseBody.outputs[0].stop_reason])); - } + } + if (resolveAIRecordingOptions().recordOutputs) { + setOutputMessages(span, getInvokeModelOutputParts(currentModelId, responseBody), finishReason); } } } @@ -425,6 +426,19 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension { } } +function getInvokeModelOutputParts(modelId: string, responseBody: ParsedChunk): GenAiMessagePart[] { + if (modelId.includes('amazon.nova')) return contentBlocksToParts(responseBody.output?.message?.content); + if (modelId.includes('anthropic.claude')) return getClaudeOutputParts(responseBody.content); + // The other families return one text: Titan, Llama, Cohere Command R, Cohere Command, Mistral. + return textParts( + responseBody.results?.[0]?.outputText ?? + responseBody.generation ?? + responseBody.text ?? + responseBody.generations?.[0]?.text ?? + responseBody.outputs?.[0]?.text, + ); +} + function getInvokeModelInputAttributes(modelId: string, requestBody: ParsedChunk): Record { if (modelId.includes('amazon.nova')) return getConverseInputAttributes(requestBody); if (modelId.includes('anthropic.claude')) return getClaudeInputAttributes(requestBody); @@ -432,7 +446,7 @@ function getInvokeModelInputAttributes(modelId: string, requestBody: ParsedChunk return getPromptInputAttributes(requestBody.inputText ?? requestBody.message ?? requestBody.prompt); } -type StreamProcessor = (parsedChunk: ParsedChunk, state: StreamResponseState) => void; +type StreamProcessor = (parsedChunk: ParsedChunk, state: BedrockStreamState) => void; function resolveStreamProcessor(modelId: string): StreamProcessor | undefined { if (modelId.includes('amazon.titan')) return processTitanChunk; @@ -450,9 +464,14 @@ async function* instrumentStream( stream: AsyncIterable, span: Span, abortSignal: AbortSignal | undefined, - processItem: (item: T, state: StreamResponseState) => void, + processItem: (item: T, state: BedrockStreamState) => void, ): AsyncGenerator { - const state: StreamResponseState = { finishReasons: [], responseTexts: [], toolCalls: [] }; + const state: BedrockStreamState = { + finishReasons: [], + responseTexts: [], + toolCalls: [], + parts: resolveAIRecordingOptions().recordOutputs ? [] : undefined, + }; try { for await (const item of stream) { processItem(item, state); @@ -462,11 +481,15 @@ async function* instrumentStream( span.setStatus({ code: SPAN_STATUS_ERROR, message: abortSignal?.aborted ? 'cancelled' : 'internal_error' }); throw error; } finally { + if (state.parts) { + setOutputMessages(span, state.parts, state.finishReasons[0]); + } endStreamSpan(span, state, false); } } -function processConverseStreamEvent(event: ConverseStreamOutput, state: StreamResponseState): void { +function processConverseStreamEvent(event: ConverseStreamOutput, state: BedrockStreamState): void { + appendConverseStreamContent(event, state); pushFinishReason(state, event.messageStop?.stopReason); const usage = event.metadata?.usage; if (usage) { @@ -481,6 +504,22 @@ function processConverseStreamEvent(event: ConverseStreamOutput, state: StreamRe } } +/** `ConverseStream` events, and Amazon Nova `InvokeModelWithResponseStream` chunks, which share their shape. */ +function appendConverseStreamContent( + { contentBlockStart, contentBlockDelta }: ParsedChunk, + state: BedrockStreamState, +): void { + const toolUse = contentBlockStart?.start?.toolUse; + if (toolUse) { + startStreamToolCall(state.parts, contentBlockStart.contentBlockIndex, toolUse.toolUseId, toolUse.name); + } + const delta = contentBlockDelta?.delta; + const index = contentBlockDelta?.contentBlockIndex; + appendStreamPart(state.parts, index, 'text', delta?.text); + appendStreamPart(state.parts, index, 'tool_call', delta?.toolUse?.input); + appendStreamPart(state.parts, index, 'reasoning', delta?.reasoningContent?.text); +} + function setStopReason(span: Span, stopReason: string | undefined): void { if (stopReason !== undefined) { span.setAttribute(GEN_AI_RESPONSE_FINISH_REASONS, JSON.stringify([stopReason])); @@ -519,7 +558,8 @@ function parseChunk(bytes?: Uint8Array): ParsedChunk { } } -function processNovaChunk(parsedChunk: ParsedChunk, state: StreamResponseState): void { +function processNovaChunk(parsedChunk: ParsedChunk, state: BedrockStreamState): void { + appendConverseStreamContent(parsedChunk, state); const usage = parsedChunk.metadata?.usage; if (typeof usage?.inputTokens === 'number') { state.promptTokens = usage.inputTokens; @@ -530,7 +570,15 @@ function processNovaChunk(parsedChunk: ParsedChunk, state: StreamResponseState): pushFinishReason(state, parsedChunk.messageStop?.stopReason); } -function processClaudeChunk(parsedChunk: ParsedChunk, state: StreamResponseState): void { +/** Claude chunks are Anthropic Messages streaming events. */ +function processClaudeChunk(parsedChunk: ParsedChunk, state: BedrockStreamState): void { + const { index, content_block: block, delta } = parsedChunk; + if (block?.type === 'tool_use') { + startStreamToolCall(state.parts, index, block.id, block.name); + } + appendStreamPart(state.parts, index, 'text', delta?.text); + appendStreamPart(state.parts, index, 'tool_call', delta?.partial_json); + appendStreamPart(state.parts, index, 'reasoning', delta?.thinking); const inputUsage = parsedChunk.message?.usage; if (typeof inputUsage?.cache_read_input_tokens === 'number') { state.cacheReadInputTokens = inputUsage.cache_read_input_tokens; @@ -550,7 +598,8 @@ function processClaudeChunk(parsedChunk: ParsedChunk, state: StreamResponseState pushFinishReason(state, parsedChunk.delta?.stop_reason); } -function processTitanChunk(parsedChunk: ParsedChunk, state: StreamResponseState): void { +function processTitanChunk(parsedChunk: ParsedChunk, state: BedrockStreamState): void { + appendStreamPart(state.parts, 0, 'text', parsedChunk.outputText); if (typeof parsedChunk.inputTextTokenCount === 'number') { state.promptTokens = parsedChunk.inputTextTokenCount; } @@ -560,7 +609,8 @@ function processTitanChunk(parsedChunk: ParsedChunk, state: StreamResponseState) pushFinishReason(state, parsedChunk.completionReason); } -function processLlamaChunk(parsedChunk: ParsedChunk, state: StreamResponseState): void { +function processLlamaChunk(parsedChunk: ParsedChunk, state: BedrockStreamState): void { + appendStreamPart(state.parts, 0, 'text', parsedChunk.generation); if (typeof parsedChunk.prompt_token_count === 'number') { state.promptTokens = parsedChunk.prompt_token_count; } @@ -584,17 +634,20 @@ function processInvocationMetrics(parsedChunk: ParsedChunk, state: StreamRespons } } -function processMistralChunk(parsedChunk: ParsedChunk, state: StreamResponseState): void { +function processMistralChunk(parsedChunk: ParsedChunk, state: BedrockStreamState): void { + appendStreamPart(state.parts, 0, 'text', parsedChunk.outputs?.[0]?.text); processInvocationMetrics(parsedChunk, state); pushFinishReason(state, parsedChunk.outputs?.[0]?.stop_reason); } -function processCohereChunk(parsedChunk: ParsedChunk, state: StreamResponseState): void { +function processCohereChunk(parsedChunk: ParsedChunk, state: BedrockStreamState): void { + appendStreamPart(state.parts, 0, 'text', parsedChunk.generations?.[0]?.text); processInvocationMetrics(parsedChunk, state); pushFinishReason(state, parsedChunk.generations?.[0]?.finish_reason); } -function processCohereRChunk(parsedChunk: ParsedChunk, state: StreamResponseState): void { +function processCohereRChunk(parsedChunk: ParsedChunk, state: BedrockStreamState): void { + appendStreamPart(state.parts, 0, 'text', parsedChunk.text); processInvocationMetrics(parsedChunk, state); pushFinishReason(state, 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 3367410fc0f0..aaca75fcfe5d 100644 --- a/packages/server-utils/test/integrations/aws-sdk-bedrock.test.ts +++ b/packages/server-utils/test/integrations/aws-sdk-bedrock.test.ts @@ -1,6 +1,7 @@ import { GEN_AI_INPUT_MESSAGES, GEN_AI_OPERATION_NAME, + GEN_AI_OUTPUT_MESSAGES, GEN_AI_PROVIDER_NAME, GEN_AI_REQUEST_MAX_TOKENS, GEN_AI_REQUEST_MODEL, @@ -557,3 +558,275 @@ describe('BedrockRuntimeServiceExtension token usage with prompt caching', () => expect(usageAttributes(span, EXPECTED_USAGE)).toEqual(EXPECTED_USAGE); }); }); + +describe('BedrockRuntimeServiceExtension output recording', () => { + const encoder = new TextEncoder(); + + function streamOf(chunks: unknown[]): AsyncIterable<{ chunk: { bytes: Uint8Array } }> { + return (async function* () { + for (const chunk of chunks) { + yield { chunk: { bytes: encoder.encode(JSON.stringify(chunk)) } }; + } + })(); + } + + function outputMessages(span: Span): unknown { + const value = spanToJSON(span).attributes?.[GEN_AI_OUTPUT_MESSAGES]; + return typeof value === 'string' ? JSON.parse(value) : value; + } + + const ASSISTANT_TURN = [ + { + role: 'assistant', + parts: [ + { type: 'reasoning', content: 'The user wants the weather.' }, + { type: 'text', content: 'Let me check.' }, + { + type: 'tool_call', + id: 'tooluse_kZJMlvQmRJ6eAyJE5GIl7Q', + name: 'get_weather', + arguments: '{"city":"Vienna"}', + }, + ], + finish_reason: 'tool_use', + }, + ]; + + beforeEach(() => { + const client = new TestClient(getDefaultTestClientOptions({ tracesSampleRate: 1 })); + setCurrentClient(client); + client.init(); + }); + + afterEach(() => { + setCurrentClient(undefined as never); + }); + + it('records the Converse response message as message parts', () => { + const span = startInactiveSpan({ name: 'bedrock' }); + + respond( + span, + 'Converse', + {}, + { + stopReason: 'tool_use', + output: { + message: { + role: 'assistant', + content: [ + { reasoningContent: { reasoningText: { text: 'The user wants the weather.', signature: 'c2ln' } } }, + { text: 'Let me check.' }, + { + toolUse: { + toolUseId: 'tooluse_kZJMlvQmRJ6eAyJE5GIl7Q', + name: 'get_weather', + input: { city: 'Vienna' }, + }, + }, + ], + }, + }, + }, + ); + + expect(outputMessages(span)).toEqual(ASSISTANT_TURN); + }); + + it('records no output message when the response has no content', () => { + const span = startInactiveSpan({ name: 'bedrock' }); + + respond(span, 'Converse', {}, { stopReason: 'end_turn', output: { message: { role: 'assistant', content: [] } } }); + + expect(spanToJSON(span).attributes).not.toHaveProperty(GEN_AI_OUTPUT_MESSAGES); + }); + + it('records the Claude InvokeModel response content as message parts', () => { + const span = startInactiveSpan({ name: 'bedrock' }); + const body = { + stop_reason: 'tool_use', + content: [ + { type: 'thinking', thinking: 'The user wants the weather.', signature: 'c2ln' }, + { type: 'redacted_thinking', data: 'ZW5jcnlwdGVk' }, + { type: 'text', text: 'Let me check.' }, + { type: 'tool_use', id: 'tooluse_kZJMlvQmRJ6eAyJE5GIl7Q', name: 'get_weather', input: { city: 'Vienna' } }, + ], + }; + + respond(span, 'InvokeModel', { modelId: CLAUDE_MODEL_ID }, { body: encoder.encode(JSON.stringify(body)) }); + + expect(outputMessages(span)).toEqual(ASSISTANT_TURN); + }); + + it.each([ + { + family: 'Titan', + modelId: 'amazon.titan-text-express-v1', + body: { results: [{ outputText: 'Paris.', completionReason: 'FINISH' }] }, + finishReason: 'FINISH', + }, + { + family: 'Llama', + modelId: 'meta.llama3-8b-instruct-v1:0', + body: { generation: 'Paris.', stop_reason: 'stop' }, + finishReason: 'stop', + }, + { + family: 'Cohere Command R', + modelId: 'cohere.command-r-v1:0', + body: { text: 'Paris.', finish_reason: 'COMPLETE' }, + finishReason: 'COMPLETE', + }, + { + family: 'Cohere Command', + modelId: 'cohere.command-text-v14', + body: { generations: [{ text: 'Paris.', finish_reason: 'COMPLETE' }] }, + finishReason: 'COMPLETE', + }, + { + family: 'Mistral', + modelId: 'mistral.mistral-7b-instruct-v0:2', + body: { outputs: [{ text: 'Paris.', stop_reason: 'stop' }] }, + finishReason: 'stop', + }, + { + family: 'Nova', + modelId: 'amazon.nova-lite-v1:0', + body: { output: { message: { role: 'assistant', content: [{ text: 'Paris.' }] } }, stopReason: 'end_turn' }, + finishReason: 'end_turn', + }, + ])('records the $family InvokeModel response text', ({ modelId, body, finishReason }) => { + const span = startInactiveSpan({ name: 'bedrock' }); + + respond(span, 'InvokeModel', { modelId }, { body: encoder.encode(JSON.stringify(body)) }); + + expect(outputMessages(span)).toEqual([ + { role: 'assistant', parts: [{ type: 'text', content: 'Paris.' }], finish_reason: finishReason }, + ]); + }); + + it('collects a ConverseStream response by content block', async () => { + const span = startInactiveSpan({ name: 'bedrock' }); + const data = { + stream: (async function* () { + yield { messageStart: { role: 'assistant' } }; + yield { contentBlockDelta: { contentBlockIndex: 0, delta: { reasoningContent: { text: 'The user wants ' } } } }; + yield { contentBlockDelta: { contentBlockIndex: 0, delta: { reasoningContent: { text: 'the weather.' } } } }; + yield { contentBlockDelta: { contentBlockIndex: 1, delta: { text: 'Let me ' } } }; + yield { contentBlockDelta: { contentBlockIndex: 1, delta: { text: 'check.' } } }; + yield { + contentBlockStart: { + contentBlockIndex: 2, + start: { toolUse: { toolUseId: 'tooluse_kZJMlvQmRJ6eAyJE5GIl7Q', name: 'get_weather' } }, + }, + }; + yield { contentBlockDelta: { contentBlockIndex: 2, delta: { toolUse: { input: '{"city":' } } } }; + yield { contentBlockDelta: { contentBlockIndex: 2, delta: { toolUse: { input: '"Vienna"}' } } } }; + yield { messageStop: { stopReason: 'tool_use' } }; + })(), + }; + + respond(span, 'ConverseStream', {}, data); + await drain(data.stream); + + expect(outputMessages(span)).toEqual(ASSISTANT_TURN); + }); + + it('collects a Claude InvokeModelWithResponseStream response by content block', async () => { + const span = startInactiveSpan({ name: 'bedrock' }); + const data = { + body: streamOf([ + { type: 'message_start', message: { role: 'assistant', usage: { input_tokens: 14, output_tokens: 1 } } }, + // Redacted reasoning has no deltas, so it leaves a gap at its index. + { type: 'content_block_start', index: 0, content_block: { type: 'redacted_thinking', data: 'ZW5jcnlwdGVk' } }, + { type: 'content_block_start', index: 1, content_block: { type: 'thinking', thinking: '' } }, + { type: 'content_block_delta', index: 1, delta: { type: 'thinking_delta', thinking: 'The user wants ' } }, + { type: 'content_block_delta', index: 1, delta: { type: 'thinking_delta', thinking: 'the weather.' } }, + { type: 'content_block_delta', index: 1, delta: { type: 'signature_delta', signature: 'c2ln' } }, + { type: 'content_block_start', index: 2, content_block: { type: 'text', text: '' } }, + { type: 'content_block_delta', index: 2, delta: { type: 'text_delta', text: 'Let me ' } }, + { type: 'content_block_delta', index: 2, delta: { type: 'text_delta', text: 'check.' } }, + { + type: 'content_block_start', + index: 3, + content_block: { type: 'tool_use', id: 'tooluse_kZJMlvQmRJ6eAyJE5GIl7Q', name: 'get_weather', input: {} }, + }, + { type: 'content_block_delta', index: 3, delta: { type: 'input_json_delta', partial_json: '{"city":' } }, + { type: 'content_block_delta', index: 3, delta: { type: 'input_json_delta', partial_json: '"Vienna"}' } }, + { type: 'message_delta', delta: { stop_reason: 'tool_use' }, usage: { output_tokens: 42 } }, + ]), + }; + + respond(span, 'InvokeModelWithResponseStream', { modelId: CLAUDE_MODEL_ID }, data); + await drain(data.body); + + expect(outputMessages(span)).toEqual(ASSISTANT_TURN); + }); + + it.each([ + { + family: 'Titan', + modelId: 'amazon.titan-text-express-v1', + toChunk: (text: string) => ({ outputText: text }), + }, + { family: 'Llama', modelId: 'meta.llama3-8b-instruct-v1:0', toChunk: (text: string) => ({ generation: text }) }, + { + family: 'Mistral', + modelId: 'mistral.mistral-large-2402-v1:0', + toChunk: (text: string) => ({ outputs: [{ text }] }), + }, + { + family: 'Cohere Command', + modelId: 'cohere.command-text-v14', + toChunk: (text: string) => ({ generations: [{ text }] }), + }, + { family: 'Cohere Command R', modelId: 'cohere.command-r-plus-v1:0', toChunk: (text: string) => ({ text }) }, + { + family: 'Nova', + modelId: 'amazon.nova-lite-v1:0', + toChunk: (text: string) => ({ contentBlockDelta: { contentBlockIndex: 0, delta: { text } } }), + }, + ])('joins the streamed $family text into one part', async ({ modelId, toChunk }) => { + const span = startInactiveSpan({ name: 'bedrock' }); + const data = { body: streamOf([toChunk('The capital'), toChunk(' is Paris.')]) }; + + respond(span, 'InvokeModelWithResponseStream', { modelId }, data); + await drain(data.body); + + expect(outputMessages(span)).toEqual([ + { role: 'assistant', parts: [{ type: 'text', content: 'The capital is Paris.' }] }, + ]); + }); + + it.each([ + { commandName: 'Converse', makeData: () => ({ output: { message: { content: [{ text: 'Paris.' }] } } }) }, + { + commandName: 'InvokeModel', + makeData: () => ({ body: encoder.encode(JSON.stringify({ outputs: [{ text: 'Paris.' }] })) }), + }, + { + commandName: 'ConverseStream', + makeData: () => ({ + stream: (async function* () { + yield { contentBlockStart: { contentBlockIndex: 0, start: { toolUse: { toolUseId: 't', name: 'f' } } } }; + yield { contentBlockDelta: { contentBlockIndex: 0, delta: { toolUse: { input: '{}' } } } }; + })(), + }), + }, + ])('records no $commandName output when dataCollection.genAI.outputs is off', async ({ commandName, makeData }) => { + const client = new TestClient( + getDefaultTestClientOptions({ tracesSampleRate: 1, dataCollection: { genAI: { outputs: false } } }), + ); + setCurrentClient(client); + client.init(); + const span = startInactiveSpan({ name: 'bedrock' }); + const data = makeData(); + + respond(span, commandName, { modelId: 'mistral.mistral-large-2402-v1:0' }, data); + if ('stream' in data) { + await drain(data.stream); + } + + expect(spanToJSON(span).attributes).not.toHaveProperty(GEN_AI_OUTPUT_MESSAGES); + }); +});