Skip to content

Commit 429ad6f

Browse files
committed
fix(llmobs): handle Bedrock stream failures without chunks
## Summary ConverseStream requests that fail before yielding data retain an empty assistant response instead of throwing from LLMObs. ## Why An `iterator.next()` rejection publishes request completion without a streamed-chunk event, so `ctx.chunks` is undefined. Direct iteration then throws while handling the SDK failure. ## Test plan `./node_modules/.bin/mocha packages/datadog-plugin-aws-sdk/test/bedrockruntime.util.spec.js` and `npm run lint`
1 parent 8de9ed2 commit 429ad6f

2 files changed

Lines changed: 56 additions & 24 deletions

File tree

packages/datadog-plugin-aws-sdk/src/services/bedrockruntime/utils.js

Lines changed: 26 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -641,7 +641,7 @@ function extractTextAndResponseReasonConverse (response) {
641641
* response, spread across start/delta chunks. We reassemble those chunks
642642
* into a normalized content-block array and reuse the non-stream extractor.
643643
*
644-
* @param {Array<object>} chunks - Ordered ConverseStreamOutput events.
644+
* @param {Array<object> | undefined} chunks - Ordered ConverseStreamOutput events.
645645
* @returns {Generation}
646646
*/
647647
function extractTextAndResponseReasonConverseFromStream (chunks) {
@@ -650,29 +650,31 @@ function extractTextAndResponseReasonConverseFromStream (chunks) {
650650
let usage = {}
651651
const blocksByIdx = new Map()
652652

653-
for (const chunk of chunks) {
654-
if (chunk.messageStart?.role) {
655-
role = chunk.messageStart.role
656-
} else if (chunk.messageStop?.stopReason) {
657-
stopReason = chunk.messageStop.stopReason
658-
} else if (chunk.metadata?.usage) {
659-
usage = chunk.metadata.usage
660-
} else if (chunk.contentBlockStart?.start?.toolUse) {
661-
const { contentBlockIndex, start: { toolUse } } = chunk.contentBlockStart
662-
blocksByIdx.set(contentBlockIndex, {
663-
toolUse: { toolUseId: toolUse.toolUseId, name: toolUse.name, inputStr: '' },
664-
})
665-
} else if (chunk.contentBlockDelta) {
666-
const { contentBlockIndex, delta } = chunk.contentBlockDelta
667-
if (typeof delta?.text === 'string') {
668-
const block = blocksByIdx.get(contentBlockIndex) ?? {}
669-
block.text = (block.text ?? '') + delta.text
670-
blocksByIdx.set(contentBlockIndex, block)
671-
} else if (typeof delta?.toolUse?.input === 'string') {
672-
const block = blocksByIdx.get(contentBlockIndex) ?? { toolUse: { inputStr: '' } }
673-
block.toolUse ??= { inputStr: '' }
674-
block.toolUse.inputStr += delta.toolUse.input
675-
blocksByIdx.set(contentBlockIndex, block)
653+
if (chunks) {
654+
for (const chunk of chunks) {
655+
if (chunk.messageStart?.role) {
656+
role = chunk.messageStart.role
657+
} else if (chunk.messageStop?.stopReason) {
658+
stopReason = chunk.messageStop.stopReason
659+
} else if (chunk.metadata?.usage) {
660+
usage = chunk.metadata.usage
661+
} else if (chunk.contentBlockStart?.start?.toolUse) {
662+
const { contentBlockIndex, start: { toolUse } } = chunk.contentBlockStart
663+
blocksByIdx.set(contentBlockIndex, {
664+
toolUse: { toolUseId: toolUse.toolUseId, name: toolUse.name, inputStr: '' },
665+
})
666+
} else if (chunk.contentBlockDelta) {
667+
const { contentBlockIndex, delta } = chunk.contentBlockDelta
668+
if (typeof delta?.text === 'string') {
669+
const block = blocksByIdx.get(contentBlockIndex) ?? {}
670+
block.text = (block.text ?? '') + delta.text
671+
blocksByIdx.set(contentBlockIndex, block)
672+
} else if (typeof delta?.toolUse?.input === 'string') {
673+
const block = blocksByIdx.get(contentBlockIndex) ?? { toolUse: { inputStr: '' } }
674+
block.toolUse ??= { inputStr: '' }
675+
block.toolUse.inputStr += delta.toolUse.input
676+
blocksByIdx.set(contentBlockIndex, block)
677+
}
676678
}
677679
}
678680
}

packages/datadog-plugin-aws-sdk/test/bedrockruntime.util.spec.js

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,36 @@ const {
99
} = require('../src/services/bedrockruntime/utils')
1010

1111
describe('bedrockruntime converse stream extractor', () => {
12+
it('returns an empty message when the stream fails before yielding a chunk', () => {
13+
const generation = extractTextAndResponseReasonConverseFromStream()
14+
15+
assert.deepStrictEqual(generation.messages, [{ role: 'assistant', content: '' }])
16+
})
17+
18+
it('aggregates streamed text, tool input, metadata, and the stop reason', () => {
19+
const generation = extractTextAndResponseReasonConverseFromStream([
20+
{ messageStart: { role: 'assistant' } },
21+
{ contentBlockDelta: { contentBlockIndex: 0, delta: { text: 'hel' } } },
22+
{ contentBlockDelta: { contentBlockIndex: 0, delta: { text: 'lo' } } },
23+
{ contentBlockDelta: { contentBlockIndex: 1, delta: { toolUse: { input: '{"city":"Berlin"}' } } } },
24+
{ metadata: { usage: { inputTokens: 2, outputTokens: 3 } } },
25+
{ messageStop: { stopReason: 'tool_use' } },
26+
])
27+
28+
assert.deepStrictEqual(generation.messages, [{
29+
role: 'assistant',
30+
content: 'hello',
31+
toolCalls: [{ name: '', arguments: { city: 'Berlin' }, toolId: '', type: 'toolUse' }],
32+
}])
33+
assert.deepStrictEqual(generation.usage, {
34+
inputTokens: 2,
35+
outputTokens: 3,
36+
cacheReadTokens: undefined,
37+
cacheWriteTokens: undefined,
38+
})
39+
assert.strictEqual(generation.finishReason, 'tool_use')
40+
})
41+
1242
it('emits empty tool-call arguments when the streamed tool-use input is malformed JSON', () => {
1343
const generation = extractTextAndResponseReasonConverseFromStream([
1444
{ messageStart: { role: 'assistant' } },

0 commit comments

Comments
 (0)