Skip to content

feat: add Bedrock InvokeModelWithResponseStream instrumentation #2845

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

yuliia-fryshko
Copy link
Contributor

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.

@yuliia-fryshko yuliia-fryshko requested a review from a team as a code owner May 22, 2025 16:25
@github-actions github-actions bot requested review from blumamir, jj22ee and trivikr May 22, 2025 16:25
@yuliia-fryshko yuliia-fryshko force-pushed the bedrock-invoke-model-stream branch 2 times, most recently from 09f9777 to eeb1e84 Compare May 23, 2025 09:56
@@ -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(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link

codecov bot commented May 27, 2025

Codecov Report

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

Project coverage is 89.73%. Comparing base (2ac08ff) to head (fcf3ff5).

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.69%   89.73%   +0.04%     
==========================================
  Files         185      185              
  Lines        9034     9093      +59     
  Branches     1852     1870      +18     
==========================================
+ Hits         8103     8160      +57     
- Misses        931      933       +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.

@yuliia-fryshko yuliia-fryshko force-pushed the bedrock-invoke-model-stream branch from eeb1e84 to 0d197f0 Compare June 2, 2025 08:36
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants