feat: add Bedrock InvokeModelWithResponseStream instrumentation #2845

yuliia-fryshko · 2025-05-22T16:25:21Z

Which problem is this PR solving?

Adds instrumentation of the InvokeModelWithResponseStreamCommand in the AWS Bedrock SDK.

Short description of the changes

Captures GenAI semantic span attributes for Amazon Titan, Anthropic Claude, Amazon Nova models.
Ensures span.end() is deferred until the full stream is consumed.
instrumentAsyncIterable is used to inspect streamed chunks in real time and extract relevant telemetry.
Includes unit tests for Titan, Claude and Nova models.

trentm · 2025-05-27T19:12:23Z

plugins/node/opentelemetry-instrumentation-aws-sdk/src/services/bedrock-runtime.ts

@@ -102,6 +102,13 @@ export class BedrockRuntimeServiceExtension implements ServiceExtension {
        return this.requestPreSpanHookConverse(request, config, diag, true);
      case 'InvokeModel':
        return this.requestPreSpanHookInvokeModel(request, config, diag);
+      case 'InvokeModelWithResponseStream':
+        return this.requestPreSpanHookInvokeModelWithResponseStream(


Is there a reason to not re-use requestPreSpanHookInvokeModel: add a 4th isStream argument and pass in false for 'InvokeModel', true for 'InvokeModelWithResponseStream', and then make the minor update to the implementation? This is how it was done for 'Converse' and 'ConverseStream'.

It looks to me like the requestPreSpanHookInvokeModel and requestPreSpanHookInvokeModelWithResponseStream functions are almost identical ... except that the latter doesn't have blocks for 'meta.llama', 'cohere.*', and 'mistral'.

Thanks for the suggestion, @trentm ! You are absolutely right! I've updated the code to consolidate requestPreSpanHookInvokeModel and requestPreSpanHookInvokeModelWithResponseStream into a single method using isStream parameter as you suggested.

codecov · 2025-05-27T19:22:05Z

Codecov Report

❌ Patch coverage is 88.09524% with 10 lines in your changes missing coverage. Please review.
✅ Project coverage is 89.81%. Comparing base (8a93515) to head (44a3dec).

Files with missing lines	Patch %	Lines
...umentation-aws-sdk/src/services/bedrock-runtime.ts	88.09%	10 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2845      +/-   ##
==========================================
- Coverage   89.83%   89.81%   -0.02%     
==========================================
  Files         188      188              
  Lines        9294     9376      +82     
  Branches     1907     1938      +31     
==========================================
+ Hits         8349     8421      +72     
- Misses        945      955      +10

Files with missing lines	Coverage Δ
...umentation-aws-sdk/src/services/bedrock-runtime.ts	`95.72% <88.09%> (-3.27%)`	⬇️

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

jj22ee · 2025-06-13T01:22:06Z

plugins/node/opentelemetry-instrumentation-aws-sdk/src/services/bedrock-runtime.ts

+  ): Promise<any> {
+    const stream = response.data?.body;
+    const modelId = response.request.commandInput?.modelId;
+    if (!stream || !span.isRecording()) return;


Suggested change

if (!stream || !span.isRecording()) return;

if (!stream) return;

!span.isRecording() is already checked before responseHookInvokeModelWithResponseStream() is called

You are right, @jj22ee! I removed this unnecessary check

jj22ee · 2025-06-13T01:38:56Z

plugins/node/opentelemetry-instrumentation-aws-sdk/src/services/bedrock-runtime.ts

+    return response.data;
+
+    // Tap into the stream at the chunk level without modifying the chunk itself.
+    function instrumentAsyncIterable<T>(


Wouldn't it be more efficient to declare these following functions outside of this member function scope?

Thank you, @jj22ee , for the review and good catches! Yes, these helper functions could be declared outside the method, at the class level, instead of being re-created every time the method is called.
Changed!

JacksonWeber · 2025-07-02T04:34:43Z

Update this to pull in the main branch and I can take another look over. Thanks!

trentm · 2025-07-02T22:25:08Z

plugins/node/opentelemetry-instrumentation-aws-sdk/src/services/bedrock-runtime.ts

+              parsedChunk,
+              span
+            );
+          }


The non-stream responseHookInvokeModel() is adding some attributes for a few more models: 'meta.llama', 'cohere.command-r', etc.
Should this function also add relevant attributes for those models? Or is it possible streaming is not supported for those models?

Also, perhaps the recordNovaAttributes(), recordClaudeAttributes(), etc. methods could be used by both responseHookInvokeModel and responseHookInvokeModelWithResponseStream. Or perhaps all the record*Attribute() methods could be moved to one setInvokeModelResponseAttributes() that is used by both responseHookInvokeModel*() methods.

trentm · 2025-07-02T22:27:58Z

plugins/node/opentelemetry-instrumentation-aws-sdk/src/services/bedrock-runtime.ts

+    // while OpenTelemetry can record span attributes from streamed data.
+    response.data.body = (async function* () {
+      try {
+        for await (const item of wrappedStream) {


IIUC, the result here is that two async iterators are created. Could the chunk handling above be moved into this for await ... and then not bother having a wrappedStream at all? I haven't tried this.

trentm · 2025-07-02T22:36:26Z

plugins/node/opentelemetry-instrumentation-aws-sdk/src/services/bedrock-runtime.ts

+      const str = Buffer.from(bytes).toString('utf-8');
+      return JSON.parse(str);
+    } catch (err) {
+      console.warn('Failed to parse streamed chunk', err);


Shouldn't use console in instrumentation code. Instead use the DiagLogger that every intsrumentation instance has.

Suggested change

console.warn('Failed to parse streamed chunk', err);

this._diag.warn('Failed to parse streamed chunk', err);

This will mean you'll need to not use static for this method.

Thank you, @trentm , for the review . I addressed all your comments and would appreciate next round of a review. Thanks

pichlermarc · 2025-07-30T16:50:06Z

@yuliia-fryshko do you have bandwidth to work on this? Looks like there are some unaddressed comments that are blocking this from getting merged.

yuliia-fryshko · 2025-08-07T14:22:20Z

@yuliia-fryshko do you have bandwidth to work on this? Looks like there are some unaddressed comments that are blocking this from getting merged.

Hi @pichlermarc ! Sorry for the long reply, I just back from the vacation today and I will try to address all comments today

…h isStream flag

…d mistral

yuliia-fryshko requested a review from a team as a code owner May 22, 2025 16:25

github-actions bot assigned blumamir and jj22ee May 22, 2025

github-actions bot added the pkg:instrumentation-aws-sdk label May 22, 2025

github-actions bot assigned trivikr May 22, 2025

github-actions bot requested review from blumamir, jj22ee and trivikr May 22, 2025 16:25

yuliia-fryshko force-pushed the bedrock-invoke-model-stream branch 2 times, most recently from 09f9777 to eeb1e84 Compare May 23, 2025 09:56

trentm reviewed May 27, 2025

View reviewed changes

yuliia-fryshko force-pushed the bedrock-invoke-model-stream branch from eeb1e84 to 0d197f0 Compare June 2, 2025 08:36

jj22ee reviewed Jun 13, 2025

View reviewed changes

yuliia-fryshko force-pushed the bedrock-invoke-model-stream branch from fe8416b to ff1cb2d Compare June 27, 2025 18:54

pichlermarc assigned trentm, trivikr, jj22ee and blumamir and unassigned trivikr, blumamir and jj22ee Jul 2, 2025

trentm reviewed Jul 2, 2025

View reviewed changes

yuliia-fryshko force-pushed the bedrock-invoke-model-stream branch from 0e2fef3 to dcd6b28 Compare August 7, 2025 14:18

yuliia-fryshko force-pushed the bedrock-invoke-model-stream branch from 1941501 to d86f003 Compare August 11, 2025 07:48

yuliia-fryshko added 2 commits August 12, 2025 23:24

feat: add Bedrock InvokeModelWithResponseStream instrumentation

e9ae8b8

refactor: unify InvokeModel pre-span hooks into a single function wit…

da2cbfa

…h isStream flag

yuliia-fryshko added 3 commits August 12, 2025 23:24

Moved static functions out of the method

0a80410

added instrumentation for Bedrock Response Stream of llama, cohere an…

ff59f72

…d mistral

added tests for cohere, cohere-r and mistral

f7bbec1

yuliia-fryshko force-pushed the bedrock-invoke-model-stream branch from 71c7a6e to f7bbec1 Compare August 12, 2025 21:24

Merge branch 'main' into bedrock-invoke-model-stream

44a3dec

	if (!stream \|\| !span.isRecording()) return;
	if (!stream) return;

	console.warn('Failed to parse streamed chunk', err);
	this._diag.warn('Failed to parse streamed chunk', err);

feat: add Bedrock InvokeModelWithResponseStream instrumentation #2845

Are you sure you want to change the base?

feat: add Bedrock InvokeModelWithResponseStream instrumentation #2845

Uh oh!

Conversation

yuliia-fryshko commented May 22, 2025

Which problem is this PR solving?

Short description of the changes

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov bot commented May 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

JacksonWeber commented Jul 2, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pichlermarc commented Jul 30, 2025

Uh oh!

yuliia-fryshko commented Aug 7, 2025

Uh oh!

Uh oh!

codecov bot commented May 27, 2025 •

edited

Loading