[MLOB-2098] feat(llmobs): record bedrock token counts

[sabrenner]

What does this PR do?

Attempts to record token counts from headers for AWS Bedrock Runtime requests (for InvokeModelCommand).

APM Patching Changes

Changes AWS patching to emit response headers on deserialization. Some clients drop headers when deserializing (like BedrockRuntime) that could be useful to us. I tried to do this a few different ways, but the most direct way seemed to be patching the deserializer for the command being executed by the client, since:

• the deserialize method is where the raw response is directly passed in before headers are stripped
• the deserialize method is not on the prototype of the Command class. While it is sometimes on the prototype of classes that inherit from Command, this changes by version and is not guaranteed for all commands. It does seem, however, that the deserialize property is available at the time a command is executed via client.send, so we patch it on the instance there.
• A small check is included to make sure deserialize is a function

LLMObs Changes

We extract the request ID and token counts to map against each other, since we don't have a shared context object being passed around. We then associate the request ID from the final response against our map, and grab the token counts from there. However, we'll first check if we were able to extract tokens from the actual response body (as some models return them) in case the header values were empty.

Additionally, changes some Bedrock shared utilities and testing fixture.

Motivation

Full support for counting tokens from models served through AWS Bedrock from @aws-sdk/client-bedrock-runtime. This is only for the LLMObs plugin but we are adding the publishers in patching.

[github-actions]

Overall package size

Self size: 8.56 MB
Deduped: 94.95 MB
No deduping: 95.47 MB

Dependency sizes

| name | version | self size | total size | |------|---------|-----------|------------| | @datadog/libdatadog | 0.4.0 | 29.44 MB | 29.44 MB | | @datadog/native-appsec | 8.4.0 | 19.25 MB | 19.26 MB | | @datadog/native-iast-taint-tracking | 3.2.0 | 13.9 MB | 13.91 MB | | @datadog/pprof | 5.5.1 | 9.79 MB | 10.17 MB | | protobufjs | 7.2.5 | 2.77 MB | 5.16 MB | | @datadog/native-iast-rewriter | 2.6.1 | 2.59 MB | 2.73 MB | | @opentelemetry/core | 1.14.0 | 872.87 kB | 1.47 MB | | @datadog/native-metrics | 3.1.0 | 1.06 MB | 1.46 MB | | @opentelemetry/api | 1.8.0 | 1.21 MB | 1.21 MB | | import-in-the-middle | 1.11.2 | 112.74 kB | 826.22 kB | | source-map | 0.7.4 | 226 kB | 226 kB | | opentracing | 0.14.7 | 194.81 kB | 194.81 kB | | lru-cache | 7.18.3 | 133.92 kB | 133.92 kB | | pprof-format | 2.1.0 | 111.69 kB | 111.69 kB | | @datadog/sketches-js | 2.1.0 | 109.9 kB | 109.9 kB | | semver | 7.6.3 | 95.82 kB | 95.82 kB | | lodash.sortby | 4.7.0 | 75.76 kB | 75.76 kB | | ignore | 5.3.1 | 51.46 kB | 51.46 kB | | shell-quote | 1.8.1 | 44.96 kB | 44.96 kB | | istanbul-lib-coverage | 3.2.0 | 29.34 kB | 29.34 kB | | rfdc | 1.3.1 | 25.21 kB | 25.21 kB | | @isaacs/ttlcache | 1.4.1 | 25.2 kB | 25.2 kB | | tlhunter-sorted-set | 0.1.0 | 24.94 kB | 24.94 kB | | limiter | 1.1.5 | 23.17 kB | 23.17 kB | | dc-polyfill | 0.1.4 | 23.1 kB | 23.1 kB | | retry | 0.13.1 | 18.85 kB | 18.85 kB | | jest-docblock | 29.7.0 | 8.99 kB | 12.76 kB | | crypto-randomuuid | 1.0.0 | 11.18 kB | 11.18 kB | | ttl-set | 1.0.0 | 4.61 kB | 9.69 kB | | path-to-regexp | 0.1.12 | 6.6 kB | 6.6 kB | | koalas | 1.0.2 | 6.47 kB | 6.47 kB | | module-details-from-path | 1.0.3 | 4.47 kB | 4.47 kB |

_{🤖 This report was automatically generated by heaviest-objects-in-the-universe}

[pr-commenter]

Benchmarks

Benchmark execution time: 2025-01-31 14:44:34

Comparing candidate commit 08d736d in PR branch sabrenner/bedrock-token-counts with baseline commit fd61c8c in branch master.

Found 0 performance improvements and 0 performance regressions! Performance is the same for 907 metrics, 26 unstable metrics.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 81.09%. Comparing base (fd61c8c) to head (08d736d).
Report is 8 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5152      +/-   ##
==========================================
+ Coverage   81.06%   81.09%   +0.03%     
==========================================
  Files         477      479       +2     
  Lines       21327    21367      +40     
==========================================
+ Hits        17288    17328      +40     
  Misses       4039     4039              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[yahya-mouman]

Great work 🚀, I'm seeing some unrelated failing tests but it should be alright.

[datadog-datadog-prod-us1]

Datadog Report

Branch report: sabrenner/bedrock-token-counts
Commit report: de2020e
Test service: dd-trace-js-integration-tests

✅ 0 Failed, 616 Passed, 0 Skipped, 15m 4.24s Total Time

[IlyasShabi]

LGTM, I'd recommend getting another approval from your team