-
Notifications
You must be signed in to change notification settings - Fork 580
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
base: main
Are you sure you want to change the base?
feat: add Bedrock InvokeModelWithResponseStream instrumentation #2845
Conversation
09f9777
to
eeb1e84
Compare
@@ -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( |
There was a problem hiding this comment.
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'.
There was a problem hiding this comment.
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.
Codecov ReportAttention: Patch coverage is
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
🚀 New features to boost your workflow:
|
eeb1e84
to
0d197f0
Compare
Which problem is this PR solving?
Adds instrumentation of the InvokeModelWithResponseStreamCommand in the AWS Bedrock SDK.
Short description of the changes
instrumentAsyncIterable
is used to inspect streamed chunks in real time and extract relevant telemetry.