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

Attention: Patch coverage is 96.72131% with 2 lines in your changes missing coverage. Please review.

Project coverage is 89.81%. Comparing base (0baef34) to head (0e2fef3).
Report is 6 commits behind head on main.

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

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2845      +/-   ##
==========================================
+ Coverage   89.76%   89.81%   +0.04%     
==========================================
  Files         187      187              
  Lines        9149     9208      +59     
  Branches     1885     1903      +18     
==========================================
+ Hits         8213     8270      +57     
- Misses        936      938       +2

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

🚀 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.

…h isStream flag

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.

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

feat: add Bedrock InvokeModelWithResponseStream instrumentation

0d197f0

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

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

fcf3ff5

…h isStream flag

jj22ee reviewed Jun 13, 2025

View reviewed changes

Moved static functions out of the method

ff1cb2d

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

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

00508e8

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

0e2fef3

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

trentm reviewed Jul 2, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat: add Bedrock InvokeModelWithResponseStream instrumentation #2845

feat: add Bedrock InvokeModelWithResponseStream instrumentation #2845

Uh oh!

yuliia-fryshko commented May 22, 2025

Uh oh!

trentm May 27, 2025

Uh oh!

yuliia-fryshko Jun 2, 2025

Uh oh!

codecov bot commented May 27, 2025 •

edited

Loading

Uh oh!

jj22ee Jun 13, 2025

Uh oh!

yuliia-fryshko Jun 27, 2025

Uh oh!

jj22ee Jun 13, 2025

Uh oh!

yuliia-fryshko Jun 27, 2025

Uh oh!

JacksonWeber commented Jul 2, 2025

Uh oh!

trentm Jul 2, 2025

Uh oh!

trentm Jul 2, 2025

Uh oh!

trentm Jul 2, 2025

Uh oh!

Uh oh!

	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!

Uh oh!

codecov bot commented May 27, 2025 •

edited

Loading